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_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 BuildType = Export,
          class CommSpaceType = Mpi>
class Halo : public CommunicationPlan<MemorySpace, CommSpaceType>
{
  public:
    using commspace_type = CommSpaceType;

    template <class IdViewType, class RankViewType, typename T = BuildType,
              std::enable_if_t<std::is_same<T, Export>::value, int> = 0>
    Halo( MPI_Comm comm, const std::size_t num_local,
          const IdViewType& element_ids, const RankViewType& element_ranks,
          const std::vector<int>& neighbor_ranks )
        : CommunicationPlan<MemorySpace, CommSpaceType>( comm )
    {
        build( num_local, element_ids, element_ranks, neighbor_ranks );
    }

    template <class IdViewType, class RankViewType, typename T = BuildType,
              std::enable_if_t<std::is_same<T, Export>::value, int> = 0>
    Halo( MPI_Comm comm, const std::size_t num_local,
          const IdViewType& element_ids, const RankViewType& element_ranks )
        : CommunicationPlan<MemorySpace, CommSpaceType>( comm )
    {
        build( num_local, element_ids, element_ranks );
    }

    template <class IdViewType, class RankViewType, typename T = BuildType,
              std::enable_if_t<std::is_same<T, Import>::value, int> = 0>
    Halo( MPI_Comm comm, const std::size_t num_local,
          const IdViewType& element_ids, const RankViewType& element_ranks,
          const std::vector<int>& neighbor_ranks )
        : CommunicationPlan<MemorySpace, CommSpaceType>( comm )
    {
        build( num_local, element_ids, element_ranks, neighbor_ranks );
    }

    template <class IdViewType, class RankViewType, typename T = BuildType,
              std::enable_if_t<std::is_same<T, Import>::value, int> = 0>
    Halo( MPI_Comm comm, const std::size_t num_local,
          const IdViewType& element_ids, const RankViewType& element_ranks )
        : CommunicationPlan<MemorySpace, CommSpaceType>( comm )
    {
        build( num_local, element_ids, element_ranks );
    }

    template <class IdViewType, class RankViewType, typename T = BuildType,
              std::enable_if_t<std::is_same<T, Export>::value, int> = 0>
    void build( const std::size_t num_local, const IdViewType& element_ids,
                const RankViewType& element_ranks,
                const std::vector<int>& neighbor_ranks )
    {
        _num_local = num_local;
        if ( element_ids.size() != element_ranks.size() )
            throw std::runtime_error( "Cabana::Halo (export): ids and ranks "
                                      "views are different sizes!" );
 
        auto neighbor_ids = this->createWithTopology(
            BuildType(), element_ranks, neighbor_ranks );
        this->createExportSteering( neighbor_ids, element_ranks, element_ids );
    }

    template <class IdViewType, class RankViewType, typename T = BuildType,
              std::enable_if_t<std::is_same<T, Export>::value, int> = 0>
    void build( const std::size_t num_local, const IdViewType& element_ids,
                const RankViewType& element_ranks )
    {
        _num_local = num_local;
        if ( element_ids.size() != element_ranks.size() )
            throw std::runtime_error( "Cabana::Halo (export): ids and ranks "
                                      "views are different sizes!" );
 
        auto neighbor_ids =
            this->createWithoutTopology( BuildType(), element_ranks );
        this->createExportSteering( neighbor_ids, element_ranks, element_ids );
    }

    template <class IdViewType, class RankViewType, typename T = BuildType,
              std::enable_if_t<std::is_same<T, Import>::value, int> = 0>
    void build( const std::size_t num_local, const IdViewType& element_ids,
                const RankViewType& element_ranks,
                const std::vector<int>& neighbor_ranks )
    {
        _num_local = num_local;
        if ( element_ids.size() != element_ranks.size() )
            throw std::runtime_error( "Cabana::Halo (import): ids and ranks "
                                      "views are different sizes!" );
 
        auto neighbor_ids_ranks_indices = this->createWithTopology(
            BuildType(), element_ranks, element_ids, neighbor_ranks );
        this->createExportSteering( std::get<0>( neighbor_ids_ranks_indices ),
                                    std::get<1>( neighbor_ids_ranks_indices ),
                                    std::get<2>( neighbor_ids_ranks_indices ) );
    }

    template <class IdViewType, class RankViewType, typename T = BuildType,
              std::enable_if_t<std::is_same<T, Import>::value, int> = 0>
    void build( const std::size_t num_local, const IdViewType& element_ids,
                const RankViewType& element_ranks )
    {
        _num_local = num_local;
        if ( element_ids.size() != element_ranks.size() )
            throw std::runtime_error( "Cabana::Halo (import): ids and ranks "
                                      "views are different sizes!" );
 
        auto neighbor_ids_ranks_indices = this->createWithoutTopology(
            BuildType(), element_ranks, element_ids );
        this->createExportSteering( std::get<0>( neighbor_ids_ranks_indices ),
                                    std::get<1>( neighbor_ids_ranks_indices ),
                                    std::get<2>( neighbor_ids_ranks_indices ) );
    }

    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, typename BuildType, typename CommSpaceType>
struct is_halo_impl<Halo<MemorySpace, BuildType, CommSpaceType>>
    : 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 commspace_type = typename HaloType::commspace_type;
    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( _comm_plan, aosoa );
    }

    auto totalSend() { return _comm_plan.totalNumExport(); }
    auto totalReceive() { return _comm_plan.totalNumImport(); }
    auto totalSend( const HaloType& halo ) { return halo.totalNumExport(); }
    auto totalReceive( const HaloType& halo ) { return halo.totalNumImport(); }

    void apply() override { applyImpl( execution_space{}, commspace_type{} ); }

    template <class ExecutionSpace, class CommSpaceType>
    std::enable_if_t<std::is_same<CommSpaceType, Mpi>::value, void>
        applyImpl( ExecutionSpace, CommSpaceType );
 
    // Future: Add applyImpl that is enabled for other CommSpaceType types.

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

    void reserve( const HaloType& halo, AoSoAType& aosoa,
                  const double overallocation )
    {
        if ( !haloCheckValidSize( halo, aosoa ) )
            throw std::runtime_error(
                "Cabana::Gather:reserve: "
                "AoSoA is the wrong size for gather! (Label: " +
                aosoa.label() + ")" );
 
        this->reserveImpl( halo, aosoa, totalSend( halo ), totalReceive( halo ),
                           overallocation );
    }
 
  private:
    using 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 commspace_type = typename HaloType::commspace_type;
    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( _comm_plan, slice );
    }

    auto totalSend() { return _comm_plan.totalNumExport(); }
    auto totalReceive() { return _comm_plan.totalNumImport(); }
    auto totalSend( const HaloType& halo ) { return halo.totalNumExport(); }
    auto totalReceive( const HaloType& halo ) { return halo.totalNumImport(); }

    void apply() override { applyImpl( execution_space{}, commspace_type{} ); }

    template <class ExecutionSpace, class CommSpaceType>
    std::enable_if_t<std::is_same<CommSpaceType, Mpi>::value, void>
        applyImpl( ExecutionSpace, CommSpaceType );
 
    // Future: Add applyImpl that is enabled for other CommSpaceType types.

    void reserve( const HaloType& halo, const SliceType& slice,
                  const double overallocation )
    {
        if ( !haloCheckValidSize( halo, slice ) )
            throw std::runtime_error(
                "Cabana::Gather:reserve: "
                "Slice is the wrong size for gather! (Label: " +
                slice.label() + ")" );
 
        // Cannot use totalSend(), totalReceive() because it may be inconsistent
        // with the new plan.
        this->reserveImpl( halo, slice, totalSend( halo ), totalReceive( halo ),
                           overallocation );
    }

    void reserve( const HaloType& halo, const SliceType& slice )
    {
        if ( !haloCheckValidSize( halo, slice ) )
            throw std::runtime_error(
                "Cabana::Gather:reserve: "
                "Slice is the wrong size for gather! (Label: " +
                slice.label() + ")" );
 
        this->reserveImpl( halo, slice, totalSend( halo ),
                           totalReceive( halo ) );
    }
 
  private:
    using 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>>
{
  public:
    static_assert( is_halo<HaloType>::value, "" );

    using commspace_type = typename HaloType::commspace_type;
    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( _comm_plan, slice );
    }

    auto totalSend() { return _comm_plan.totalNumImport(); }
    auto totalReceive() { return _comm_plan.totalNumExport(); }
    auto totalSend( const HaloType& halo ) { return halo.totalNumImport(); }
    auto totalReceive( const HaloType& halo ) { return halo.totalNumExport(); }

    void apply() override { applyImpl( execution_space{}, commspace_type{} ); }

    template <class ExecutionSpace, class CommSpaceType>
    std::enable_if_t<std::is_same<CommSpaceType, Mpi>::value, void>
        applyImpl( ExecutionSpace, CommSpaceType );
 
    // Future: Add applyImpl that is enabled for other CommSpaceType types.

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

    void reserve( const HaloType& halo, const SliceType& slice )
    {
        if ( !haloCheckValidSize( halo, slice ) )
            throw std::runtime_error(
                "Cabana::Scatter::reserve: "
                "Slice is the wrong size for scatter! (Label: " +
                slice.label() + ")" );
 
        this->reserveImpl( halo, slice, totalSend( halo ),
                           totalReceive( halo ) );
    }
 
  private:
    using base_type::_comm_plan;
    using base_type::_recv_size;
    using base_type::_send_size;
};
 
} // end namespace Cabana
 
// Include communication backends from what is enabled in CMake.
#ifdef Cabana_ENABLE_MPI
#include <impl/Cabana_Halo_Mpi.hpp>
#endif // Enable MPI
 
namespace Cabana
{

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