Cabana_Halo.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_HALO_HPP
#define CABANA_HALO_HPP
 
#include <Cabana_AoSoA.hpp>
#include <Cabana_CommunicationPlan.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 Halo : public CommunicationPlan<MemorySpace>
{
  public:
    template <class IdViewType, class RankViewType>
    Halo( MPI_Comm comm, const std::size_t num_local,
          const IdViewType& element_export_ids,
          const RankViewType& element_export_ranks,
          const std::vector<int>& neighbor_ranks )
        : CommunicationPlan<MemorySpace>( comm )
        , _num_local( num_local )
    {
        if ( element_export_ids.size() != element_export_ranks.size() )
            throw std::runtime_error( "Export ids and ranks different sizes!" );
 
        auto neighbor_ids = this->createFromExportsAndTopology(
            element_export_ranks, neighbor_ranks );
        this->createExportSteering( neighbor_ids, element_export_ranks,
                                    element_export_ids );
    }

    template <class IdViewType, class RankViewType>
    Halo( MPI_Comm comm, const std::size_t num_local,
          const IdViewType& element_export_ids,
          const RankViewType& element_export_ranks )
        : CommunicationPlan<MemorySpace>( comm )
        , _num_local( num_local )
    {
        if ( element_export_ids.size() != element_export_ranks.size() )
            throw std::runtime_error( "Export ids and ranks different sizes!" );
 
        auto neighbor_ids = this->createFromExportsOnly( element_export_ranks );
        this->createExportSteering( neighbor_ids, element_export_ranks,
                                    element_export_ids );
    }

    std::size_t numLocal() const { return _num_local; }

    std::size_t numGhost() const { return this->totalNumImport(); }
 
  private:
    std::size_t _num_local;
};
 
//---------------------------------------------------------------------------//
template <typename>
struct is_halo_impl : public std::false_type
{
};
 
template <typename MemorySpace>
struct is_halo_impl<Halo<MemorySpace>> : public std::true_type
{
};

template <class T>
struct is_halo : public is_halo_impl<typename std::remove_cv<T>::type>::type
{
};

template <class Halo, class ParticleData>
bool haloCheckValidSize(
    const Halo& halo, const ParticleData& particles,
    typename std::enable_if<( is_halo<Halo>::value ), int>::type* = 0 )
{
    // Check that the data is the right size.
    return ( particles.size() == halo.numLocal() + halo.numGhost() );
}
 
template <class HaloType, class AoSoAType, class SFINAE = void>
class Gather;
 
//---------------------------------------------------------------------------//
template <class HaloType, class AoSoAType>
class Gather<HaloType, AoSoAType,
             typename std::enable_if<is_aosoa<AoSoAType>::value>::type>
    : public CommunicationData<HaloType, CommunicationDataAoSoA<AoSoAType>>
{
  public:
    static_assert( is_halo<HaloType>::value, "" );

    using base_type =
        CommunicationData<HaloType, CommunicationDataAoSoA<AoSoAType>>;
    using plan_type = typename base_type::plan_type;
    using execution_space = typename base_type::execution_space;
    using memory_space = typename base_type::memory_space;
    using data_type = typename base_type::data_type;
    using buffer_type = typename base_type::buffer_type;

    Gather( HaloType halo, AoSoAType aosoa, const double overallocation = 1.0 )
        : base_type( halo, aosoa, overallocation )
    {
        reserve( _halo, aosoa );
    }

    auto totalSend() { return _halo.totalNumExport(); }
    auto totalReceive() { return _halo.totalNumImport(); }

    template <class ExecutionSpace>
    void apply( ExecutionSpace )
    {
        Kokkos::Profiling::ScopedRegion region( "Cabana::gather" );
 
        // Get the buffers and particle data (local copies for lambdas below).
        auto send_buffer = this->getSendBuffer();
        auto recv_buffer = this->getReceiveBuffer();
        auto aosoa = this->getData();
 
        // Get the steering vector for the sends.
        auto steering = _halo.getExportSteering();
        // Gather from the local data into a tuple-contiguous send buffer.
        auto gather_send_buffer_func = KOKKOS_LAMBDA( const std::size_t i )
        {
            send_buffer( i ) = aosoa.getTuple( steering( i ) );
        };
        Kokkos::RangePolicy<ExecutionSpace> send_policy( 0, _send_size );
        Kokkos::parallel_for( "Cabana::gather::gather_send_buffer", send_policy,
                              gather_send_buffer_func );
        Kokkos::fence();
 
        // The halo has it's own communication space so choose any mpi tag.
        const int mpi_tag = 2345;
 
        // Post non-blocking receives.
        int num_n = _halo.numNeighbor();
        std::vector<MPI_Request> requests( num_n );
        std::pair<std::size_t, std::size_t> recv_range = { 0, 0 };
        for ( int n = 0; n < num_n; ++n )
        {
            recv_range.second = recv_range.first + _halo.numImport( n );
 
            auto recv_subview = Kokkos::subview( recv_buffer, recv_range );
 
            MPI_Irecv( recv_subview.data(),
                       recv_subview.size() * sizeof( data_type ), MPI_BYTE,
                       _halo.neighborRank( n ), mpi_tag, _halo.comm(),
                       &( requests[n] ) );
 
            recv_range.first = recv_range.second;
        }
 
        // Do blocking sends.
        std::pair<std::size_t, std::size_t> send_range = { 0, 0 };
        for ( int n = 0; n < num_n; ++n )
        {
            send_range.second = send_range.first + _halo.numExport( n );
 
            auto send_subview = Kokkos::subview( send_buffer, send_range );
 
            MPI_Send( send_subview.data(),
                      send_subview.size() * sizeof( data_type ), MPI_BYTE,
                      _halo.neighborRank( n ), mpi_tag, _halo.comm() );
 
            send_range.first = send_range.second;
        }
 
        // Wait on non-blocking receives.
        std::vector<MPI_Status> status( num_n );
        const int ec =
            MPI_Waitall( requests.size(), requests.data(), status.data() );
        if ( MPI_SUCCESS != ec )
            throw std::logic_error( "Failed MPI Communication" );
 
        // Extract the receive buffer into the ghosted elements.
        std::size_t num_local = _halo.numLocal();
        auto extract_recv_buffer_func = KOKKOS_LAMBDA( const std::size_t i )
        {
            std::size_t ghost_idx = i + num_local;
            aosoa.setTuple( ghost_idx, recv_buffer( i ) );
        };
        Kokkos::RangePolicy<ExecutionSpace> recv_policy( 0, _recv_size );
        Kokkos::parallel_for( "Cabana::gather::extract_recv_buffer",
                              recv_policy, extract_recv_buffer_func );
        Kokkos::fence();
 
        // Barrier before completing to ensure synchronization.
        MPI_Barrier( _halo.comm() );
    }
 
    void apply() override { apply( execution_space{} ); }

    void reserve( const HaloType& halo, AoSoAType& aosoa )
    {
        if ( !haloCheckValidSize( halo, aosoa ) )
            throw std::runtime_error( "AoSoA is the wrong size for gather!" );
 
        this->reserveImpl( halo, aosoa, totalSend(), totalReceive() );
    }

    void reserve( const HaloType& halo, AoSoAType& aosoa,
                  const double overallocation )
    {
        if ( !haloCheckValidSize( halo, aosoa ) )
            throw std::runtime_error( "AoSoA is the wrong size for gather!" );
 
        this->reserveImpl( halo, aosoa, totalSend(), totalReceive(),
                           overallocation );
    }
 
  private:
    plan_type _halo = base_type::_comm_plan;
    using base_type::_recv_size;
    using base_type::_send_size;
};

template <class HaloType, class SliceType>
class Gather<HaloType, SliceType,
             typename std::enable_if<is_slice<SliceType>::value>::type>
    : public CommunicationData<HaloType, CommunicationDataSlice<SliceType>>
{
  public:
    static_assert( is_halo<HaloType>::value, "" );

    using base_type =
        CommunicationData<HaloType, CommunicationDataSlice<SliceType>>;
    using plan_type = typename base_type::plan_type;
    using execution_space = typename base_type::execution_space;
    using memory_space = typename base_type::memory_space;
    using data_type = typename base_type::data_type;
    using buffer_type = typename base_type::buffer_type;

    Gather( HaloType halo, SliceType slice, const double overallocation = 1.0 )
        : base_type( halo, slice, overallocation )
    {
        reserve( _halo, slice );
    }

    auto totalSend() { return _halo.totalNumExport(); }
    auto totalReceive() { return _halo.totalNumImport(); }

    template <class ExecutionSpace>
    void apply( ExecutionSpace )
    {
        Kokkos::Profiling::ScopedRegion region( "Cabana::gather" );
 
        // Get the buffers (local copies for lambdas below).
        auto send_buffer = this->getSendBuffer();
        auto recv_buffer = this->getReceiveBuffer();
        auto slice = this->getData();
 
        // Get the number of components in the slice.
        std::size_t num_comp = this->getSliceComponents();
 
        // Get the raw slice data.
        auto slice_data = slice.data();
 
        // Get the steering vector for the sends.
        auto steering = _halo.getExportSteering();
 
        // Gather from the local data into a tuple-contiguous send buffer.
        auto gather_send_buffer_func = KOKKOS_LAMBDA( const std::size_t i )
        {
            auto s = SliceType::index_type::s( steering( i ) );
            auto a = SliceType::index_type::a( steering( i ) );
            std::size_t slice_offset = s * slice.stride( 0 ) + a;
            for ( std::size_t n = 0; n < num_comp; ++n )
                send_buffer( i, n ) =
                    slice_data[slice_offset + n * SliceType::vector_length];
        };
        Kokkos::RangePolicy<ExecutionSpace> send_policy( 0, _send_size );
        Kokkos::parallel_for( "Cabana::gather::gather_send_buffer", send_policy,
                              gather_send_buffer_func );
        Kokkos::fence();
 
        // The halo has it's own communication space so choose any mpi tag.
        const int mpi_tag = 2345;
 
        // Post non-blocking receives.
        int num_n = _halo.numNeighbor();
        std::vector<MPI_Request> requests( num_n );
        std::pair<std::size_t, std::size_t> recv_range = { 0, 0 };
        for ( int n = 0; n < num_n; ++n )
        {
            recv_range.second = recv_range.first + _halo.numImport( n );
 
            auto recv_subview =
                Kokkos::subview( recv_buffer, recv_range, Kokkos::ALL );
 
            MPI_Irecv( recv_subview.data(),
                       recv_subview.size() * sizeof( data_type ), MPI_BYTE,
                       _halo.neighborRank( n ), mpi_tag, _halo.comm(),
                       &( requests[n] ) );
 
            recv_range.first = recv_range.second;
        }
 
        // Do blocking sends.
        std::pair<std::size_t, std::size_t> send_range = { 0, 0 };
        for ( int n = 0; n < num_n; ++n )
        {
            send_range.second = send_range.first + _halo.numExport( n );
 
            auto send_subview =
                Kokkos::subview( send_buffer, send_range, Kokkos::ALL );
 
            MPI_Send( send_subview.data(),
                      send_subview.size() * sizeof( data_type ), MPI_BYTE,
                      _halo.neighborRank( n ), mpi_tag, _halo.comm() );
 
            send_range.first = send_range.second;
        }
 
        // Wait on non-blocking receives.
        std::vector<MPI_Status> status( num_n );
        const int ec =
            MPI_Waitall( requests.size(), requests.data(), status.data() );
        if ( MPI_SUCCESS != ec )
            throw std::logic_error( "Failed MPI Communication" );
 
        // Extract the receive buffer into the ghosted elements.
        std::size_t num_local = _halo.numLocal();
        auto extract_recv_buffer_func = KOKKOS_LAMBDA( const std::size_t i )
        {
            std::size_t ghost_idx = i + num_local;
            auto s = SliceType::index_type::s( ghost_idx );
            auto a = SliceType::index_type::a( ghost_idx );
            std::size_t slice_offset = s * slice.stride( 0 ) + a;
            for ( std::size_t n = 0; n < num_comp; ++n )
                slice_data[slice_offset + SliceType::vector_length * n] =
                    recv_buffer( i, n );
        };
        Kokkos::RangePolicy<ExecutionSpace> recv_policy( 0, _recv_size );
        Kokkos::parallel_for( "Cabana::gather::extract_recv_buffer",
                              recv_policy, extract_recv_buffer_func );
        Kokkos::fence();
 
        // Barrier before completing to ensure synchronization.
        MPI_Barrier( _halo.comm() );
    }
 
    void apply() override { apply( execution_space{} ); }

    void reserve( const HaloType& halo, const SliceType& slice,
                  const double overallocation )
    {
        if ( !haloCheckValidSize( halo, slice ) )
            throw std::runtime_error( "AoSoA is the wrong size for gather!" );
 
        this->reserveImpl( halo, slice, totalSend(), totalReceive(),
                           overallocation );
    }

    void reserve( const HaloType& halo, const SliceType& slice )
    {
        if ( !haloCheckValidSize( halo, slice ) )
            throw std::runtime_error( "AoSoA is the wrong size for gather!" );
 
        this->reserveImpl( halo, slice, totalSend(), totalReceive() );
    }
 
  private:
    plan_type _halo = base_type::_comm_plan;
    using base_type::_recv_size;
    using base_type::_send_size;
};
 
//---------------------------------------------------------------------------//
template <class HaloType, class ParticleDataType>
auto createGather( const HaloType& halo, const ParticleDataType& data,
                   const double overallocation = 1.0 )
{
    return Gather<HaloType, ParticleDataType>( halo, data, overallocation );
}
 
//---------------------------------------------------------------------------//
template <class HaloType, class ParticleDataType>
void gather( const HaloType& halo, ParticleDataType& data )
{
    auto gather = createGather( halo, data );
    gather.apply();
}
 
/**********
 * SCATTER *
 **********/
 
//---------------------------------------------------------------------------//
template <class HaloType, class SliceType>
class Scatter
    : public CommunicationData<HaloType, CommunicationDataSlice<SliceType>>
{
    static_assert( is_halo<HaloType>::value, "" );
 
  public:
    using base_type =
        CommunicationData<HaloType, CommunicationDataSlice<SliceType>>;
    using plan_type = typename base_type::plan_type;
    using execution_space = typename base_type::execution_space;
    using memory_space = typename base_type::memory_space;
    using data_type = typename base_type::data_type;
    using buffer_type = typename base_type::buffer_type;

    Scatter( HaloType halo, SliceType slice, const double overallocation = 1.0 )
        : base_type( halo, slice, overallocation )
    {
        reserve( _halo, slice );
    }

    auto totalSend() { return _halo.totalNumImport(); }
    auto totalReceive() { return _halo.totalNumExport(); }

    template <class ExecutionSpace>
    void apply( ExecutionSpace )
    {
        Kokkos::Profiling::ScopedRegion region( "Cabana::scatter" );
 
        // Get the buffers (local copies for lambdas below).
        auto send_buffer = this->getSendBuffer();
        auto recv_buffer = this->getReceiveBuffer();
        auto slice = this->getData();
 
        // Get the number of components in the slice.
        std::size_t num_comp = this->getSliceComponents();
 
        // Get the raw slice data. Wrap in a 1D Kokkos View so we can unroll the
        // components of each slice element.
        Kokkos::View<data_type*, memory_space,
                     Kokkos::MemoryTraits<Kokkos::Unmanaged>>
            slice_data( slice.data(), slice.numSoA() * slice.stride( 0 ) );
 
        // Extract the send buffer from the ghosted elements.
        std::size_t num_local = _halo.numLocal();
        auto extract_send_buffer_func = KOKKOS_LAMBDA( const std::size_t i )
        {
            std::size_t ghost_idx = i + num_local;
            auto s = SliceType::index_type::s( ghost_idx );
            auto a = SliceType::index_type::a( ghost_idx );
            std::size_t slice_offset = s * slice.stride( 0 ) + a;
            for ( std::size_t n = 0; n < num_comp; ++n )
                send_buffer( i, n ) =
                    slice_data( slice_offset + SliceType::vector_length * n );
        };
        Kokkos::RangePolicy<ExecutionSpace> send_policy( 0, _send_size );
        Kokkos::parallel_for( "Cabana::scatter::extract_send_buffer",
                              send_policy, extract_send_buffer_func );
        Kokkos::fence();
 
        // The halo has it's own communication space so choose any mpi tag.
        const int mpi_tag = 2345;
 
        // Post non-blocking receives.
        int num_n = _halo.numNeighbor();
        std::vector<MPI_Request> requests( num_n );
        std::pair<std::size_t, std::size_t> recv_range = { 0, 0 };
        for ( int n = 0; n < num_n; ++n )
        {
            recv_range.second = recv_range.first + _halo.numExport( n );
 
            auto recv_subview =
                Kokkos::subview( recv_buffer, recv_range, Kokkos::ALL );
 
            MPI_Irecv( recv_subview.data(),
                       recv_subview.size() * sizeof( data_type ), MPI_BYTE,
                       _halo.neighborRank( n ), mpi_tag, _halo.comm(),
                       &( requests[n] ) );
 
            recv_range.first = recv_range.second;
        }
 
        // Do blocking sends.
        std::pair<std::size_t, std::size_t> send_range = { 0, 0 };
        for ( int n = 0; n < num_n; ++n )
        {
            send_range.second = send_range.first + _halo.numImport( n );
 
            auto send_subview =
                Kokkos::subview( send_buffer, send_range, Kokkos::ALL );
 
            MPI_Send( send_subview.data(),
                      send_subview.size() * sizeof( data_type ), MPI_BYTE,
                      _halo.neighborRank( n ), mpi_tag, _halo.comm() );
 
            send_range.first = send_range.second;
        }
 
        // Wait on non-blocking receives.
        std::vector<MPI_Status> status( num_n );
        const int ec =
            MPI_Waitall( requests.size(), requests.data(), status.data() );
        if ( MPI_SUCCESS != ec )
            throw std::logic_error( "Failed MPI Communication" );
 
        // Get the steering vector for the sends.
        auto steering = _halo.getExportSteering();
 
        // Scatter the ghosts in the receive buffer into the local values.
        auto scatter_recv_buffer_func = KOKKOS_LAMBDA( const std::size_t i )
        {
            auto s = SliceType::index_type::s( steering( i ) );
            auto a = SliceType::index_type::a( steering( i ) );
            std::size_t slice_offset = s * slice.stride( 0 ) + a;
            for ( std::size_t n = 0; n < num_comp; ++n )
                Kokkos::atomic_add(
                    &slice_data( slice_offset + SliceType::vector_length * n ),
                    recv_buffer( i, n ) );
        };
        Kokkos::RangePolicy<ExecutionSpace> recv_policy( 0, _recv_size );
        Kokkos::parallel_for( "Cabana::scatter::scatter_recv_buffer",
                              recv_policy, scatter_recv_buffer_func );
        Kokkos::fence();
 
        // Barrier before completing to ensure synchronization.
        MPI_Barrier( _halo.comm() );
    }
 
    void apply() override { apply( execution_space{} ); }

    void reserve( const HaloType& halo, const SliceType& slice,
                  const double overallocation )
    {
        if ( !haloCheckValidSize( halo, slice ) )
            throw std::runtime_error( "AoSoA is the wrong size for scatter!" );
 
        this->reserveImpl( halo, slice, totalSend(), totalReceive(),
                           overallocation );
    }

    void reserve( const HaloType& halo, const SliceType& slice )
    {
        if ( !haloCheckValidSize( halo, slice ) )
            throw std::runtime_error( "AoSoA is the wrong size for scatter!" );
 
        this->reserveImpl( halo, slice, totalSend(), totalReceive() );
    }
 
  private:
    plan_type _halo = base_type::_comm_plan;
    using base_type::_recv_size;
    using base_type::_send_size;
};

template <class HaloType, class SliceType>
auto createScatter( const HaloType& halo, const SliceType& slice,
                    const double overallocation = 1.0,
                    typename std::enable_if<( is_halo<HaloType>::value &&
                                              is_slice<SliceType>::value ),
                                            int>::type* = 0 )
{
    return Scatter<HaloType, SliceType>( halo, slice, overallocation );
}
 
//---------------------------------------------------------------------------//
template <class HaloType, class SliceType>
void scatter( const HaloType& halo, SliceType& slice,
              typename std::enable_if<( is_halo<HaloType>::value &&
                                        is_slice<SliceType>::value ),
                                      int>::type* = 0 )
{
    auto scatter = createScatter( halo, slice );
    scatter.apply();
}
 
//---------------------------------------------------------------------------//
 
} // end namespace Cabana
 
#endif // end CABANA_HALO_HPP