mpi_kron.h

Go to the documentation of this file.

#pragma once
#include <cassert>
#include <vector>
#include <algorithm> // std::copy_n
#include <map>
#include "exceptions.h"
 
 
namespace dg
{
namespace detail{
// I think the mpi registry should not be exposed to the User
// It may be an idea to hint at it in the grid documentation
// inline declared functions generate no warning about "defined but not used"...
// and can be defined in multiple translation units
struct MPICartInfo
{
    MPI_Comm root;
    std::vector<int> remain_dims;
    // std::vector<int> periods; // no need to track because the root and all the children have the same periods
};
//we keep track of communicators that were created in the past
// inline variables have the same address in multiple translation units
// and have external linkage by default
inline std::map<MPI_Comm, MPICartInfo> mpi_cart_registry;
 
inline void mpi_cart_registry_display( std::ostream& out = std::cout)
{
    out << "MPI Cart registry\n";
    for ( auto pair : detail::mpi_cart_registry)
    {
        auto& re = pair.second.remain_dims;
        out << "Comm "<<pair.first<<" Root "<<pair.second.root<<" Remains size "<<re.size();
        out << " r:";
        for( uint i=0; i<re.size(); i++)
            out << " "<<re[i];
        out << std::endl;
    }
}
 
inline void mpi_cart_registry_clear( )
{
    mpi_cart_registry.clear();
}
 
inline void mpi_cart_register_cart( MPI_Comm comm)
{
    // Register Cartesian comm if not already there
    auto it = detail::mpi_cart_registry.find( comm);
    if( it == detail::mpi_cart_registry.end())
    {
        int ndims;
        MPI_Cartdim_get( comm, &ndims);
        std::vector<int> remain_dims( ndims, true);
        detail::MPICartInfo info{ comm, remain_dims};
        detail::mpi_cart_registry[comm] = info;
    }
}
}
 
inline void register_mpi_cart_sub( MPI_Comm comm, const int remain_dims[], MPI_Comm newcomm)
{
    detail::mpi_cart_register_cart( comm);
    detail::MPICartInfo info = detail::mpi_cart_registry.at(comm);
    for( unsigned u=0; u<info.remain_dims.size(); u++)
        info.remain_dims[u]  = remain_dims[u];
    detail::mpi_cart_registry[newcomm] = info;
}
// Can't decide whether to make register_mpi_cart_sub public ...
 
inline MPI_Comm mpi_cart_sub( MPI_Comm comm, std::vector<int> remain_dims, bool
    duplicate = false)
{
    int ndims;
    MPI_Cartdim_get( comm, &ndims);
    assert( (unsigned) ndims == remain_dims.size());
 
    detail::mpi_cart_register_cart( comm);
    detail::MPICartInfo info = detail::mpi_cart_registry.at(comm);
    // info.remain_dims may be larger than remain_dims because comm may have a parent
    // but exactly remain_dims.size() entries are true
    int counter =0;
    for( unsigned u=0; u<info.remain_dims.size(); u++)
        if( info.remain_dims[u])
        {
            info.remain_dims[u]  = remain_dims[counter];
            counter ++;
        }
    assert( counter == (int)remain_dims.size());
    if( not duplicate)
    {
        for (auto it = detail::mpi_cart_registry.begin(); it !=
            detail::mpi_cart_registry.end(); ++it)
        {
//            int comp_root;
//            MPI_Comm_compare( it->second.root, info.root, &comp_root);
//            if( (comp_root == MPI_IDENT or comp_root == MPI_CONGRUENT) and
//                it->second.remain_dims == info.remain_dims)
            if( it->second.root == info.root && it->second.remain_dims ==
                info.remain_dims)
            {
                return it->first;
            }
        }
    }
    MPI_Comm newcomm;
    int err = MPI_Cart_sub( comm, &remain_dims[0], &newcomm);
    if( err != MPI_SUCCESS)
        throw Error(Message(_ping_)<<
                "Cannot create Cartesian sub comm from given communicator");
    register_mpi_cart_sub( comm, &remain_dims[0], newcomm);
    return newcomm;
}
 
inline MPI_Comm mpi_cart_kron( std::vector<MPI_Comm> comms)
{
    // This non-template interface must exist so compiler can deduce call
    // mpi_cart_kron( {comm0, comm1});
    if ( comms.empty())
        return MPI_COMM_NULL;
    if( comms.size() == 1)
        return comms[0];
 
    std::vector<detail::MPICartInfo> infos(comms.size());
 
    for( unsigned u=0; u<comms.size(); u++)
    {
    try{
        infos [u] = detail::mpi_cart_registry.at(comms[u]);
    }catch( std::exception& e)
    {
        std::cerr << "Did not find "<<comms[u]<<" in Registry!\n";
        detail::mpi_cart_registry_display();
        throw e;
 
    }
    }
    MPI_Comm root = infos[0].root;
    for( unsigned u=0; u<comms.size(); u++)
        if( infos[u].root != root)
            throw Error(Message(_ping_)<<
                    "In mpi_cart_kron all comms must have same root comm "
                    <<root<<" Offending comm number "<<u<<" with root "
                    <<infos[u].root);
    auto root_info = detail::mpi_cart_registry.at(root);
    size_t ndims = root_info.remain_dims.size();
    std::vector<int> remain_dims( ndims, false) ;
    unsigned current_free_k=0;
    for( unsigned u=0; u<comms.size(); u++)
    {
        for( unsigned k=0; k<ndims; k++)
        if( infos[u].remain_dims[k])
        {
            if( remain_dims[k])
                throw Error(Message(_ping_)<<
                    "Cannot form kronecker product with given communicators as remaining dims must be orthogonal ");
            if( k < current_free_k)
                throw Error(Message(_ping_)<<
                    "Cannot form kronecker product with communicators in reverse order compared to original ");
            remain_dims[k] = infos[u].remain_dims[k];
            current_free_k = k+1;
        }
    }
    return dg::mpi_cart_sub( root_info.root, remain_dims, false);
}
 
template<class Vector>
MPI_Comm mpi_cart_kron( Vector comms)
{
    return mpi_cart_kron( std::vector<MPI_Comm>(comms.begin(), comms.end()));
}
 
 
inline std::vector<MPI_Comm> mpi_cart_split( MPI_Comm comm)
{
    // Check that there is a Comm that was already split
    int ndims;
    MPI_Cartdim_get( comm, &ndims);
 
    std::vector<MPI_Comm> comms(ndims);
    for( int u=0; u<ndims; u++)
    {
        std::vector<int> remain_dims(ndims, 0);
        remain_dims[u] = 1;
        comms[u] = mpi_cart_sub( comm, remain_dims, false);
    }
    return comms;
}
template<size_t Nd>
std::array<MPI_Comm, Nd> mpi_cart_split_as( MPI_Comm comm)
{
    auto split = mpi_cart_split( comm);
    std::array<MPI_Comm, Nd> arr;
    std::copy_n( split.begin(), Nd, arr.begin());
    return arr;
}
 
 
 
// Need to think about those again
// /*! @brief unregister a communicator
// * For example if a communicator was freed
// * @param comm delete associated registry entry
// */
//void unregister_mpi_comm( MPI_Comm comm)
//{
//    detail::mpi_cart_registry.erase( comm);
//}
//
// /*! @brief call \c MPI_Comm_free(comm) followed by \c dg::unregister_mpi_comm(comm)
// * @param comm free communicator and delete associated registry entry
// */
//void mpi_comm_free( MPI_Comm * comm)
//{
//    MPI_Comm_free(comm);
//    unregister_mpi_comm( *comm);
//}
 
 
 
}