gmsh_mpi.h

#pragma once

#include <vector>
#include <map>
#include <array>
#include <mpi/mpi.h>

#include "eigen.h"

#if ((!defined(NDEBUG)) && defined(USE_MPI))
    #define ASSERT_MPI_RANK_0 {                 \
        int __rank;                             \
        MPI_Comm_rank(MPI_COMM_WORLD, &__rank); \
        assert(__rank == 0);                    \
    }
#else
    #define ASSERT_MPI_RANK_0 (void) 0
#endif

template<typename T>
void
priv_MPI_Send(Eigen::Matrix<T, Eigen::Dynamic, 1>& vec, int dst, MPI_Comm comm)
{
    Eigen::Index vsize = vec.size();
    MPI_Send(&vsize, sizeof(Eigen::Index), MPI_PACKED, dst, 0, comm);
    MPI_Send(vec.data(), vec.size()*sizeof(T), MPI_PACKED, dst, 0, comm);
}

template<typename T>
void
priv_MPI_Recv(Eigen::Matrix<T, Eigen::Dynamic, 1>& vec, int src, MPI_Comm comm)
{
    MPI_Status status;
    Eigen::Index vsize;
    MPI_Recv(&vsize, sizeof(Eigen::Index), MPI_PACKED, src, 0, comm, &status);
    vec.resize(vsize);
    MPI_Recv(vec.data(), vec.size()*sizeof(T), MPI_PACKED, src, 0, comm, &status);
}

template<typename T>
void
priv_MPI_Send(Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>& mat,
    int dst, MPI_Comm comm)
{
    std::array<Eigen::Index, 2> msizes;
    msizes[0] = mat.rows();
    msizes[1] = mat.cols();
    MPI_Send(&msizes, sizeof(msizes), MPI_PACKED, dst, 0, comm);
    MPI_Send(mat.data(), mat.rows()*mat.cols()*sizeof(T), MPI_PACKED, dst, 0, comm);
}

template<typename T>
void
priv_MPI_Recv(Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>& mat,
    int src, MPI_Comm comm)
{
    MPI_Status status;
    std::array<Eigen::Index, 2> msizes;
    MPI_Recv(&msizes, sizeof(msizes), MPI_PACKED, src, 0, comm, &status);
    mat.resize(msizes[0], msizes[1]);
    MPI_Recv(mat.data(), mat.rows()*mat.cols()*sizeof(T), MPI_PACKED, src, 0, comm, &status);
}

template<typename T>
void
priv_MPI_Send(std::vector<T>& vec, int dst, MPI_Comm comm)
{
    static_assert(std::is_trivially_copyable<T>::value, "Type must be trivially copyable");
    size_t vsize = vec.size();
    MPI_Send(&vsize, 1, MPI_UNSIGNED_LONG_LONG, dst, 0, comm);
    MPI_Send(vec.data(), vec.size()*sizeof(T), MPI_PACKED, dst, 0, comm);
}

template<typename T>
void
priv_MPI_Recv(std::vector<T>& vec, int src, MPI_Comm comm)
{
    static_assert(std::is_trivially_copyable<T>::value, "Type must be trivially copyable");
    MPI_Status status;
    size_t vsize;
    MPI_Recv(&vsize, 1, MPI_UNSIGNED_LONG_LONG, src, 0, comm, &status);
    vec.resize(vsize);
    MPI_Recv(vec.data(), vec.size()*sizeof(T), MPI_PACKED, src, 0, comm, &status);
}



template<typename T>
void
priv_MPI_Bcast(std::vector<T>& vec, int root, MPI_Comm comm)
{
    static_assert(std::is_trivially_copyable<T>::value, "Type must be trivially copyable");

    int rank;
    MPI_Comm_rank(comm, &rank);

    if (rank == root)
    {
        size_t vsize = vec.size();
        MPI_Bcast(&vsize, 1, MPI_UNSIGNED_LONG_LONG, root, comm);
        MPI_Bcast(vec.data(), vec.size()*sizeof(T), MPI_PACKED, root, comm);
    }
    else
    {
        size_t vsize;
        MPI_Bcast(&vsize, 1, MPI_UNSIGNED_LONG_LONG, root, comm);
        vec.resize(vsize);
        MPI_Bcast(vec.data(), vec.size()*sizeof(T), MPI_PACKED, root, comm);
    }
}

template<typename T1, typename T2>
void
priv_MPI_Bcast(std::map<T1, std::vector<T2>>& map, int root, MPI_Comm comm)
{
    static_assert(std::is_trivially_copyable<T1>::value, "Type must be trivially copyable");
    static_assert(std::is_trivially_copyable<T2>::value, "Type must be trivially copyable");

    int rank;
    MPI_Comm_rank(comm, &rank);

    if (rank == root)
    {
        size_t msize = map.size();
        MPI_Bcast(&msize, 1, MPI_UNSIGNED_LONG_LONG, root, comm);
        for (auto& [l, rv] : map)
        {
            auto ll = l;
            MPI_Bcast(&ll, sizeof(T1), MPI_PACKED, root, comm);
            size_t vsize = rv.size();
            MPI_Bcast(&vsize, 1, MPI_UNSIGNED_LONG_LONG, root, comm);
            MPI_Bcast(rv.data(), vsize*sizeof(T2), MPI_PACKED, root, comm);
        }
    }
    else
    {
        size_t msize;
        MPI_Bcast(&msize, 1, MPI_UNSIGNED_LONG_LONG, root, comm);
        for (size_t i = 0; i < msize; i++)
        {
            T1 l;
            MPI_Bcast(&l, sizeof(T1), MPI_PACKED, root, comm);
            size_t vsize;
            MPI_Bcast(&vsize, 1, MPI_UNSIGNED_LONG_LONG, root, comm);
            std::vector<T2> rv;
            rv.resize(vsize);
            MPI_Bcast(rv.data(), vsize*sizeof(T2), MPI_PACKED, root, comm);
            map[l] = std::move(rv);
        }
    }
}