my_mpi.hpp

Go to the documentation of this file.

#ifndef MY_MPI_HPP
#define MY_MPI_HPP
 
#include <mpi.h>
#include <vector>
#include <cstdio>
 
extern MPI_Comm SML_COMM_WORLD;
extern int SML_COMM_RANK;
extern int SML_COMM_SIZE;
 
void create_comm_world(int color);
void destroy_comm_world();
 
/* To illustrate the different communicators: Imagine SML_COMM_WORLD has 12 ranks:
 * 0  1  2  3  4  5  6  7  8  9 10 11
 */
 
struct MyMPI{
    // General communicator
    MPI_Comm comm;
    MPI_Group group;
    int nranks;
    int my_rank;
    /* Reorders SML_COMM_WORLD e.g. (with 12 ranks, 4 planes):
     * 0   3   6   9
     * 1   4   7  10
     * 2   5   8  11
     * So e.g. processes with ranks 0 and 3 in SML_COMM_WORLD are reordered to be
     * consecutive in comm
     * Not sure why this transpose is done; maybe the order affects some MPI algorithms
     * but not clear when/why this would be better.
     * Also, in the Fortran code comm overwrites SML_COMM_WORLD in the setup (both
     * called sml_comm), so something similar could be done here. - ALS
     */
 
    // Intraplane communicator
    MPI_Comm plane_comm;
    int my_plane_rank;
    int n_plane_ranks;
    /* Takes a column from comm, e.g. (with 12 ranks, 4 planes, this is rank 4):
     *     3
     *     4
     *     5
     */
 
    // Interplane communicator
    MPI_Comm intpl_comm;
    int my_intpl_rank;
    int n_intpl_ranks;
    /* Takes a row from comm, e.g. (with 12 ranks, 4 planes, this is rank 4):
     *
     * 1   4   7  10
     *
     */
 
    MyMPI(){}
 
    // Spoofed MPI for dry run
    MyMPI(int nranks_in)
      : nranks(nranks_in),
        my_rank(0) {}
 
    MyMPI(const MPI_Comm& comm_world, int nplanes);
};
 
// Some useful debugging tools
template <typename F>
void execute_in_rank_order(const MPI_Comm& comm, F func){
    int my_rank;
    int n_ranks;
    MPI_Comm_rank( comm, &my_rank );
    MPI_Comm_size( comm, &n_ranks );
 
    // double to make sure *every* rank is flushing at *every* barrier
    fflush(stdout);
    MPI_Barrier(comm);
    fflush(stdout);
    MPI_Barrier(comm);
    for(int i=0; i<n_ranks; i++){
        if(my_rank==i) func(i);
        fflush(stdout);
        MPI_Barrier(comm);
        fflush(stdout);
        MPI_Barrier(comm);
    }
};
 
template <typename T>
MPI_Datatype get_mpi_type() {
    if constexpr(std::is_same<T, double>::value) {
        return MPI_DOUBLE;
    } else if constexpr(std::is_same<T, int>::value) {
        return MPI_INT;
    } else if constexpr(std::is_same<T, float>::value) {
        return MPI_FLOAT;
    } else if constexpr(std::is_same<T, long long int>::value) {
        return MPI_LONG_LONG_INT;
    } else {
        // Fallback for unsupported types
        static_assert(!std::is_same<T, T>::value, "Unsupported MPI data type");
        return MPI_DATATYPE_NULL;
    }
}
 
// Function overload for Kokkos Views
template <typename ViewType>
MPI_Datatype get_mpi_type(const ViewType& view) {
    using value_type = typename ViewType::value_type;
    return get_mpi_type<value_type>();
}
 
 
#endif