Cabana_CommunicationPlanBase.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_COMMUNICATIONPLANBASE_HPP
#define CABANA_COMMUNICATIONPLANBASE_HPP
 
#include <Cabana_Core_Config.hpp>
#include <Cabana_Tags.hpp>
#include <Cabana_Utils.hpp>
 
#include <Kokkos_Core.hpp>
#include <Kokkos_ScatterView.hpp>
 
#include <mpi.h>
 
#include <algorithm>
#include <exception>
#include <memory>
#include <numeric>
#include <type_traits>
#include <vector>
 
namespace Cabana
{
//---------------------------------------------------------------------------//
 
namespace Impl
{
//---------------------------------------------------------------------------//
// Count sends and create steering algorithm tags.
struct CountSendsAndCreateSteeringDuplicated
{
};
struct CountSendsAndCreateSteeringAtomic
{
};
 
//---------------------------------------------------------------------------//
// Count sends and create steering algorithm selector.
template <class ExecutionSpace>
struct CountSendsAndCreateSteeringAlgorithm;
 
// CUDA and HIP use atomics.
#ifdef KOKKOS_ENABLE_CUDA
template <>
struct CountSendsAndCreateSteeringAlgorithm<Kokkos::Cuda>
{
    using type = CountSendsAndCreateSteeringAtomic;
};
#endif // end KOKKOS_ENABLE_CUDA
#ifdef KOKKOS_ENABLE_HIP
template <>
struct CountSendsAndCreateSteeringAlgorithm<Kokkos::HIP>
{
    using type = CountSendsAndCreateSteeringAtomic;
};
#endif // end KOKKOS_ENABLE_HIP
#ifdef KOKKOS_ENABLE_SYCL
template <>
struct CountSendsAndCreateSteeringAlgorithm<Kokkos::Experimental::SYCL>
{
    using type = CountSendsAndCreateSteeringAtomic;
};
#endif // end KOKKOS_ENABLE_SYCL
#ifdef KOKKOS_ENABLE_OPENMPTARGET
template <>
struct CountSendsAndCreateSteeringAlgorithm<Kokkos::Experimental::OpenMPTarget>
{
    using type = CountSendsAndCreateSteeringAtomic;
};
#endif // end KOKKOS_ENABLE_OPENMPTARGET
 
// The default is to use duplication.
template <class ExecutionSpace>
struct CountSendsAndCreateSteeringAlgorithm
{
    using type = CountSendsAndCreateSteeringDuplicated;
};
 
//---------------------------------------------------------------------------//
// Count sends and generate the steering vector. Atomic version.
template <class ExecutionSpace, class ExportRankView>
auto countSendsAndCreateSteering( ExecutionSpace,
                                  const ExportRankView element_export_ranks,
                                  const int comm_size,
                                  CountSendsAndCreateSteeringAtomic )
    -> std::pair<Kokkos::View<int*, typename ExportRankView::memory_space>,
                 Kokkos::View<typename ExportRankView::size_type*,
                              typename ExportRankView::memory_space>>
{
    using memory_space = typename ExportRankView::memory_space;
    using size_type = typename ExportRankView::size_type;
 
    // Create views.
    Kokkos::View<int*, memory_space> neighbor_counts( "neighbor_counts",
                                                      comm_size );
    Kokkos::View<size_type*, memory_space> neighbor_ids(
        Kokkos::ViewAllocateWithoutInitializing( "neighbor_ids" ),
        element_export_ranks.size() );
 
    // Count the sends and create the steering vector.
 
    // For smaller values of comm_size, use the optimized scratch memory path
    // Value of 64 is arbitrary, should be at least 27 to cover 3-d halos?
    if ( comm_size <= 64 )
    {
        constexpr int team_size = 256;
        Kokkos::TeamPolicy<ExecutionSpace> team_policy(
            ( element_export_ranks.size() + team_size - 1 ) / team_size,
            team_size );
        team_policy = team_policy.set_scratch_size(
            0,
            Kokkos::PerTeam( sizeof( int ) * ( team_size + 2 * comm_size ) ) );
        Kokkos::parallel_for(
            "Cabana::CommunicationPlan::countSendsAndCreateSteeringShared",
            team_policy,
            KOKKOS_LAMBDA(
                const typename Kokkos::TeamPolicy<ExecutionSpace>::member_type&
                    team ) {
                // NOTE: `get_shmem` returns shared memory pointers *aligned to
                // 8 bytes*, so if `comm_size` is odd we can get erroneously
                // padded offsets if we call `get_shmem` separately for each
                // shared memory intermediary. Acquiring all the needed scratch
                // memory at once then computing pointer offsets by hand avoids
                // this issue.
                int* scratch = (int*)team.team_shmem().get_shmem(
                    ( team.team_size() + 2 * comm_size ) * sizeof( int ), 0 );
 
                // local neighbor_ids, gives the local offset relative to
                // calculated global offset. Size team.team_size() * sizeof(int)
                int* local_neighbor_ids = scratch;
                // local histogram, `comm_size` in size
                int* histo = local_neighbor_ids + team.team_size();
                // offset into global array, `comm_size` in size
                int* global_offset = histo + comm_size;
                // overall element `tid`
                const auto tid =
                    team.team_rank() + team.league_rank() * team.team_size();
                // local element
                const int local_id = team.team_rank();
                // total number of elements, for convenience
                const int num_elements = element_export_ranks.size();
                // my element export rank
                const int my_element_export_rank =
                    ( tid < num_elements ? element_export_ranks( tid ) : -1 );
 
                // cannot outright terminate early b/c threads share work
                // if (tid >= num_elements) return;
                const bool in_bounds =
                    tid < num_elements && my_element_export_rank >= 0;
 
                Kokkos::parallel_for(
                    Kokkos::TeamThreadRange( team, comm_size ),
                    [&]( const int i ) { histo[i] = 0; } );
 
                // synchronize zeroing
                team.team_barrier();
 
                // build local histogram, need to encode num_elements check here
                if ( in_bounds )
                {
                    // shared memory atomic add, accumulate into local offset
                    local_neighbor_ids[local_id] = Kokkos::atomic_fetch_add(
                        &histo[my_element_export_rank], 1 );
                }
 
                // synchronize local histogram build
                team.team_barrier();
 
                // reserve space in global array via a loop over neighbor counts
                Kokkos::parallel_for(
                    Kokkos::TeamThreadRange( team, comm_size ),
                    [&]( const int i )
                    {
                        // global memory atomic add, reserves space
                        global_offset[i] = Kokkos::atomic_fetch_add(
                            &neighbor_counts( i ), histo[i] );
                    } );
 
                // synchronize block-stride loop
                team.team_barrier();
 
                // and now store to my location
                if ( in_bounds )
                {
                    neighbor_ids( tid ) =
                        global_offset[my_element_export_rank] +
                        local_neighbor_ids[local_id];
                }
            } );
    }
    else
    {
        // For larger numbers of export ranks (for ex, migration) we use the
        // global memory version, though this is a point of future optimization
 
        Kokkos::parallel_for(
            "Cabana::CommunicationPlan::countSendsAndCreateSteering",
            Kokkos::RangePolicy<ExecutionSpace>( 0,
                                                 element_export_ranks.size() ),
            KOKKOS_LAMBDA( const size_type i ) {
                if ( element_export_ranks( i ) >= 0 )
                    neighbor_ids( i ) = Kokkos::atomic_fetch_add(
                        &neighbor_counts( element_export_ranks( i ) ), 1 );
            } );
    }
    Kokkos::fence();
 
    // Return the counts and ids.
    return std::make_pair( neighbor_counts, neighbor_ids );
}
 
//---------------------------------------------------------------------------//
// Count sends and generate the steering vector. Duplicated version.
template <class ExecutionSpace, class ExportRankView>
auto countSendsAndCreateSteering( ExecutionSpace,
                                  const ExportRankView element_export_ranks,
                                  const int comm_size,
                                  CountSendsAndCreateSteeringDuplicated )
    -> std::pair<Kokkos::View<int*, typename ExportRankView::memory_space>,
                 Kokkos::View<typename ExportRankView::size_type*,
                              typename ExportRankView::memory_space>>
{
    using memory_space = typename ExportRankView::memory_space;
    using size_type = typename ExportRankView::size_type;
 
    // Create a unique thread token.
    Kokkos::Experimental::UniqueToken<
        ExecutionSpace, Kokkos::Experimental::UniqueTokenScope::Global>
        unique_token;
 
    // Create views.
    Kokkos::View<int*, memory_space> neighbor_counts(
        Kokkos::ViewAllocateWithoutInitializing( "neighbor_counts" ),
        comm_size );
    Kokkos::View<size_type*, memory_space> neighbor_ids(
        Kokkos::ViewAllocateWithoutInitializing( "neighbor_ids" ),
        element_export_ranks.size() );
    Kokkos::View<int**, memory_space> neighbor_counts_dup(
        "neighbor_counts", unique_token.size(), comm_size );
    Kokkos::View<size_type**, memory_space> neighbor_ids_dup(
        "neighbor_ids", unique_token.size(), element_export_ranks.size() );
 
    // Compute initial duplicated sends and steering.
    Kokkos::parallel_for(
        "Cabana::CommunicationPlan::intialCount",
        Kokkos::RangePolicy<ExecutionSpace>( 0, element_export_ranks.size() ),
        KOKKOS_LAMBDA( const size_type i ) {
            if ( element_export_ranks( i ) >= 0 )
            {
                // Get the thread id.
                auto thread_id = unique_token.acquire();
 
                // Do the duplicated fetch-add. If this is a valid element id
                // increment the send count for this rank. Add the incremented
                // count as the thread-local neighbor id. This is too big by
                // one (because we use the prefix increment operator) but we
                // want a non-zero value so we can later find which thread
                // this element was located on because we are always
                // guaranteed a non-zero value. We will subtract this value
                // later.
                neighbor_ids_dup( thread_id, i ) = ++neighbor_counts_dup(
                    thread_id, element_export_ranks( i ) );
 
                // Release the thread id.
                unique_token.release( thread_id );
            }
        } );
    Kokkos::fence();
 
    // Team policy
    using team_policy =
        Kokkos::TeamPolicy<ExecutionSpace, Kokkos::Schedule<Kokkos::Dynamic>>;
    using index_type = typename team_policy::index_type;
 
    // Compute the send counts for each neighbor rank by reducing across
    // the thread duplicates.
    Kokkos::parallel_for(
        "Cabana::CommunicationPlan::finalCount",
        team_policy( neighbor_counts.extent( 0 ), Kokkos::AUTO ),
        KOKKOS_LAMBDA( const typename team_policy::member_type& team ) {
            // Get the element id.
            auto i = team.league_rank();
 
            // Add the thread results.
            int thread_counts = 0;
            Kokkos::parallel_reduce(
                Kokkos::TeamThreadRange( team,
                                         neighbor_counts_dup.extent( 0 ) ),
                [&]( const index_type thread_id, int& result )
                { result += neighbor_counts_dup( thread_id, i ); },
                thread_counts );
            neighbor_counts( i ) = thread_counts;
        } );
    Kokkos::fence();
 
    // Compute the location of each export element in the send buffer of
    // its destination rank.
    Kokkos::parallel_for(
        "Cabana::CommunicationPlan::createSteering",
        team_policy( element_export_ranks.size(), Kokkos::AUTO ),
        KOKKOS_LAMBDA( const typename team_policy::member_type& team ) {
            // Get the element id.
            auto i = team.league_rank();
 
            // Only operate on valid elements
            if ( element_export_ranks( i ) >= 0 )
            {
                // Compute the thread id in which we located the element
                // during the count phase. Only the thread in which we
                // located the element will contribute to the reduction.
                index_type dup_thread = 0;
                Kokkos::parallel_reduce(
                    Kokkos::TeamThreadRange( team,
                                             neighbor_ids_dup.extent( 0 ) ),
                    [&]( const index_type thread_id, index_type& result )
                    {
                        if ( neighbor_ids_dup( thread_id, i ) > 0 )
                            result += thread_id;
                    },
                    dup_thread );
 
                // Compute the offset of this element in the steering
                // vector for its destination rank. Loop through the
                // threads up to the thread that found this element in the
                // count stage. All thread counts prior to that thread
                // will contribute to the offset.
                size_type thread_offset = 0;
                Kokkos::parallel_reduce(
                    Kokkos::TeamThreadRange( team, dup_thread ),
                    [&]( const index_type thread_id, size_type& result ) {
                        result += neighbor_counts_dup(
                            thread_id, element_export_ranks( i ) );
                    },
                    thread_offset );
 
                // Add the thread-local value to the offset where we subtract
                // the 1 that we added artificially when we were first
                // counting.
                neighbor_ids( i ) =
                    thread_offset + neighbor_ids_dup( dup_thread, i ) - 1;
            }
        } );
    Kokkos::fence();
 
    // Return the counts and ids.
    return std::make_pair( neighbor_counts, neighbor_ids );
}
 
//---------------------------------------------------------------------------//
} // end namespace Impl
 
//---------------------------------------------------------------------------//
inline std::vector<int> getUniqueTopology( MPI_Comm comm,
                                           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( comm, &my_rank );
    for ( auto& n : topology )
    {
        if ( n == my_rank )
        {
            std::swap( n, topology[0] );
            break;
        }
    }
    return topology;
}
 
//---------------------------------------------------------------------------//
template <class MemorySpace>
class CommunicationPlanBase
{
  public:
    using memory_space = MemorySpace;
    static_assert( Kokkos::is_memory_space<MemorySpace>() );

    using execution_space = typename memory_space::execution_space;

    using size_type = typename memory_space::memory_space::size_type;
 
  protected:
    CommunicationPlanBase( MPI_Comm comm )
    {
        _comm_ptr.reset(
            // Duplicate the communicator and store in a std::shared_ptr so that
            // all copies point to the same object
            [comm]()
            {
                auto p = std::make_unique<MPI_Comm>();
                MPI_Comm_dup( comm, p.get() );
                return p.release();
            }(),
            // Custom deleter to mark the communicator for deallocation
            []( MPI_Comm* p )
            {
                MPI_Comm_free( p );
                delete p;
            } );
    }
 
  public:
    MPI_Comm comm() const { return *_comm_ptr; }

    int numNeighbor() const { return _neighbors.size(); }

    int neighborRank( const int neighbor ) const
    {
        return _neighbors[neighbor];
    }

    std::size_t numExport( const int neighbor ) const
    {
        return _num_export[neighbor];
    }

    std::size_t totalNumExport() const { return _total_num_export; }

    std::size_t numImport( const int neighbor ) const
    {
        return _num_import[neighbor];
    }

    std::size_t totalNumImport() const { return _total_num_import; }

    std::size_t exportSize() const { return _num_export_element; }

    Kokkos::View<std::size_t*, memory_space> getExportSteering() const
    {
        return _export_steering;
    }
 
    // The functions in the public block below would normally be protected but
    // we make them public to allow using private class data in CUDA kernels
    // with lambda functions.
  public:
    template <class PackViewType, class RankViewType>
    void createExportSteering( const PackViewType& neighbor_ids,
                               const RankViewType& element_export_ranks )
    {
        // passing in element_export_ranks here as a dummy argument.
        createSteering( true, neighbor_ids, element_export_ranks,
                        element_export_ranks );
    }

    template <class PackViewType, class RankViewType, class IdViewType>
    void createExportSteering( const PackViewType& neighbor_ids,
                               const RankViewType& element_export_ranks,
                               const IdViewType& element_export_ids )
    {
        createSteering( false, neighbor_ids, element_export_ranks,
                        element_export_ids );
    }

    // Create the export steering vector.
    template <class ExecutionSpace, class PackViewType, class RankViewType,
              class IdViewType>
    void createSteering( ExecutionSpace, const bool use_iota,
                         const PackViewType& neighbor_ids,
                         const RankViewType& element_export_ranks,
                         const IdViewType& element_export_ids )
    {
        static_assert( is_accessible_from<memory_space, ExecutionSpace>{}, "" );
 
        if ( !use_iota &&
             ( element_export_ids.size() != element_export_ranks.size() ) )
            throw std::runtime_error(
                "Cabana::CommunicationPlan::createSteering: Export ids and "
                "ranks different sizes!" );
 
        // Get the size of this communicator.
        int comm_size = -1;
        MPI_Comm_size( *_comm_ptr, &comm_size );
 
        // Calculate the steering offsets via exclusive prefix sum for the
        // exports.
        int num_n = _neighbors.size();
        std::vector<std::size_t> offsets( num_n, 0.0 );
        for ( int n = 1; n < num_n; ++n )
            offsets[n] = offsets[n - 1] + _num_export[n - 1];
 
        // Map the offsets to the device.
        Kokkos::View<std::size_t*, Kokkos::HostSpace> rank_offsets_host(
            Kokkos::ViewAllocateWithoutInitializing( "rank_map" ), comm_size );
        for ( int n = 0; n < num_n; ++n )
            rank_offsets_host( _neighbors[n] ) = offsets[n];
        auto rank_offsets = Kokkos::create_mirror_view_and_copy(
            memory_space(), rank_offsets_host );
 
        // Create the export steering vector for writing local elements into
        // the send buffer. Note we create a local, shallow copy - this is a
        // CUDA workaround for handling class private data.
        _export_steering = Kokkos::View<std::size_t*, memory_space>(
            Kokkos::ViewAllocateWithoutInitializing( "export_steering" ),
            _total_num_export );
        auto steer_vec = _export_steering;
        Kokkos::parallel_for(
            "Cabana::CommunicationPlan::createSteering",
            Kokkos::RangePolicy<ExecutionSpace>( 0, _num_export_element ),
            KOKKOS_LAMBDA( const int i ) {
                if ( element_export_ranks( i ) >= 0 )
                    steer_vec( rank_offsets( element_export_ranks( i ) ) +
                               neighbor_ids( i ) ) =
                        ( use_iota ) ? i : element_export_ids( i );
            } );
        Kokkos::fence();
    }
 
    template <class PackViewType, class RankViewType, class IdViewType>
    void createSteering( const bool use_iota, const PackViewType& neighbor_ids,
                         const RankViewType& element_export_ranks,
                         const IdViewType& element_export_ids )
    {
        // Use the default execution space.
        createSteering( execution_space{}, use_iota, neighbor_ids,
                        element_export_ranks, element_export_ids );
    }
 
  protected:
    std::shared_ptr<MPI_Comm> _comm_ptr;
    std::vector<int> _neighbors;
    std::size_t _total_num_export;
    std::size_t _total_num_import;
    std::vector<std::size_t> _num_export;
    std::vector<std::size_t> _num_import;
    std::size_t _num_export_element;
    Kokkos::View<std::size_t*, memory_space> _export_steering;
};
 
//---------------------------------------------------------------------------//
template <class AoSoAType>
struct CommunicationDataAoSoA
{
    static_assert( is_aosoa<AoSoAType>::value, "" );

    using particle_data_type = AoSoAType;
    using memory_space = typename particle_data_type::memory_space;
    using data_type = typename particle_data_type::tuple_type;
    using buffer_type = typename Kokkos::View<data_type*, memory_space>;

    CommunicationDataAoSoA( particle_data_type particles )
        : _particles( particles )
    {
        _send_buffer = buffer_type(
            Kokkos::ViewAllocateWithoutInitializing( "send_buffer" ), 0 );
        _recv_buffer = buffer_type(
            Kokkos::ViewAllocateWithoutInitializing( "recv_buffer" ), 0 );
    }

    void reallocateSend( const std::size_t num_send )
    {
        Kokkos::realloc( _send_buffer, num_send );
    }

    void reallocateReceive( const std::size_t num_recv )
    {
        Kokkos::realloc( _recv_buffer, num_recv );
    }

    buffer_type _send_buffer;
    buffer_type _recv_buffer;
    particle_data_type _particles;
    std::size_t _num_comp = 0;
};

template <class SliceType>
struct CommunicationDataSlice
{
    static_assert( is_slice<SliceType>::value, "" );

    using particle_data_type = SliceType;
    using memory_space = typename particle_data_type::memory_space;
    using data_type = typename particle_data_type::value_type;
    using buffer_type =
        typename Kokkos::View<data_type**, Kokkos::LayoutRight, memory_space>;

    CommunicationDataSlice( particle_data_type particles )
        : _particles( particles )
    {
        setSliceComponents();
 
        _send_buffer = buffer_type(
            Kokkos::ViewAllocateWithoutInitializing( "send_buffer" ), 0, 0 );
        _recv_buffer = buffer_type(
            Kokkos::ViewAllocateWithoutInitializing( "recv_buffer" ), 0, 0 );
    }

    void reallocateSend( const std::size_t num_send )
    {
        Kokkos::realloc( _send_buffer, num_send, _num_comp );
    }

    void reallocateReceive( const std::size_t num_recv )
    {
        Kokkos::realloc( _recv_buffer, num_recv, _num_comp );
    }

    void setSliceComponents()
    {
        _num_comp = 1;
        for ( std::size_t d = 2; d < _particles.viewRank(); ++d )
            _num_comp *= _particles.extent( d );
    }

    buffer_type _send_buffer;
    buffer_type _recv_buffer;
    particle_data_type _particles;
    std::size_t _num_comp;
};
//---------------------------------------------------------------------------//

template <class CommPlanType, class CommDataType>
class CommunicationDataBase
{
  public:
    using plan_type = CommPlanType;
    using execution_space = typename plan_type::execution_space;
    using policy_type = Kokkos::RangePolicy<execution_space>;
    using comm_data_type = CommDataType;
    using particle_data_type = typename comm_data_type::particle_data_type;
    using memory_space = typename comm_data_type::memory_space;
    using data_type = typename comm_data_type::data_type;
    using buffer_type = typename comm_data_type::buffer_type;
 
  protected:
    CommunicationDataBase( const CommPlanType& comm_plan,
                           const particle_data_type& particles,
                           const double overallocation = 1.0 )
        : _comm_plan( comm_plan )
        , _comm_data( CommDataType( particles ) )
        , _overallocation( overallocation )
    {
    }
 
  public:
    buffer_type getSendBuffer() const { return _comm_data._send_buffer; }
    buffer_type getReceiveBuffer() const { return _comm_data._recv_buffer; }

    particle_data_type getData() const { return _comm_data._particles; }
    void setData( const particle_data_type& particles )
    {
        _comm_data._particles = particles;
    }

    auto sendSize() { return _send_size; }
    auto receiveSize() { return _recv_size; }
    auto sendCapacity() { return _comm_data._send_buffer.extent( 0 ); }
    auto receiveCapacity() { return _comm_data._recv_buffer.extent( 0 ); }
    void shrinkToFit( const bool use_overallocation = false )
    {
        auto shrunk_send_size = _send_size;
        auto shrunk_recv_size = _recv_size;
        if ( use_overallocation )
        {
            shrunk_send_size *= _overallocation;
            shrunk_recv_size *= _overallocation;
        }
        _comm_data.reallocateSend( shrunk_send_size );
        _comm_data.reallocateReceive( shrunk_recv_size );
    }

    virtual void apply() = 0;

    template <class ExecutionSpace>
    void apply( ExecutionSpace );

    void reserveImpl( const CommPlanType& comm_plan,
                      const particle_data_type particles,
                      const std::size_t total_send,
                      const std::size_t total_recv,
                      const double overallocation )
    {
        if ( overallocation < 1.0 )
            throw std::runtime_error( "Cabana::CommunicationPlan: "
                                      "Cannot allocate buffers with less space "
                                      "than data to communicate!" );
        _overallocation = overallocation;
 
        reserveImpl( comm_plan, particles, total_send, total_recv );
    }
    void reserveImpl( const CommPlanType& comm_plan,
                      const particle_data_type particles,
                      const std::size_t total_send,
                      const std::size_t total_recv )
    {
        _comm_plan = comm_plan;
        setData( particles );
 
        auto send_capacity = sendCapacity();
        auto new_send_size = static_cast<std::size_t>(
            static_cast<double>( total_send ) * _overallocation );
        if ( new_send_size > send_capacity )
            _comm_data.reallocateSend( new_send_size );
 
        auto recv_capacity = receiveCapacity();
        auto new_recv_size = static_cast<std::size_t>(
            static_cast<double>( total_recv ) * _overallocation );
        if ( new_recv_size > recv_capacity )
            _comm_data.reallocateReceive( new_recv_size );
 
        _send_size = total_send;
        _recv_size = total_recv;
    }
 
  protected:
    auto getSliceComponents() { return _comm_data._num_comp; };

    plan_type _comm_plan;
    comm_data_type _comm_data;
    double _overallocation;
    std::size_t _send_size;
    std::size_t _recv_size;
};
 
// Forward declaration of the primary CommunicationPlan template.
template <class MemorySpace, class CommSpaceType = Mpi>
class CommunicationPlan;
 
// Forward declaration of the primary CommunicationData template.
template <class CommPlanType, class CommDataType, class CommSpaceType = Mpi>
class CommunicationData;
 
} // namespace Cabana
 
// Include communication backends from what is enabled in CMake.
#ifdef Cabana_ENABLE_MPI
#include <impl/Cabana_CommunicationPlan_Mpi.hpp>
#endif // Enable MPI
 
#endif // end CABANA_COMMUNICATIONPLANBASE_HPP