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){
        // Get total number of processors
        MPI_Comm_size( comm_world, &nranks );

        // Check that nranks is divisible by nplanes
        if(nranks%nplanes != 0){
            fflush(stdout);
            if(SML_COMM_RANK==0){
                printf("\nERROR: Invalid number of MPI ranks: must be divisible by sml_nphi_total\n");
                printf("#MPI ranks=%d", nranks);
                printf("#sml_nphi_total=%d", nplanes);
                fflush(stdout);
            }
            MPI_Abort(SML_COMM_WORLD, 1);
        }

        int ranks_per_plane = nranks/nplanes;

        if(false/*sml_plane_major*/){
            // get the group underlying sml_comm
            //mpi_comm_group(sml_comm, sml_comm_group, ierr)
        }else{
            // redefine sml_comm pe ordering from consecutive within planes
            // (Plane_Major) to consecutive across planes (Interplane-Major)
            std::vector<int> new_sml_comm_ranks(nranks);
            int k = 0;
            for(int j=0; j<ranks_per_plane; j++){
                for(int i=0; i<nplanes; i++){
                    new_sml_comm_ranks[ranks_per_plane*i + j] = k;
                    k++;
                }
            }

            // get the group underlying sml_comm (== mpi_comm_world)
            // Use sml_comm_world instead of MPI_COMM_WORLD (for MPMD/staging execution)
            MPI_Group mpi_comm_world_group;
            MPI_Comm_group(comm_world, &mpi_comm_world_group);

            // create new group permuting sml_comm pe ordering to Interplane-Major
            MPI_Group_incl(mpi_comm_world_group, nranks, new_sml_comm_ranks.data(), &group);

            // Create the new communicator
            MPI_Comm_create(comm_world, group, &comm);
            MPI_Comm_rank( comm, &my_rank );
        }
        
        // INTRA-PLANE MPI COMMUNICATOR
        //checkpoint("\nPlane mpi communication");
        int plane_0_pe=int(my_rank/ranks_per_plane)*ranks_per_plane;
        std::vector<int> plane_ranks(ranks_per_plane);
        for(int j=0; j<ranks_per_plane; j++){
            plane_ranks[j]=plane_0_pe + j;
        }

        // Create the new plane group
        MPI_Group plane_group;
        MPI_Group_incl(group, ranks_per_plane, plane_ranks.data(), &plane_group);

        // Create the new plane communicator
        MPI_Comm_create(comm, plane_group, &plane_comm);

        //call mpi_comm_size(plane_comm, plane_nranks, ierr)
        MPI_Comm_rank(plane_comm, &my_plane_rank);
        MPI_Comm_size(plane_comm, &n_plane_ranks);

        // INTER-PLANE MPI COMMUNICATOR
        //checkpoint("\nInter-plane mpi communication");
        std::vector<int> intpl_ranks(nplanes);
        for(int i=0; i<nplanes; i++){
            intpl_ranks[i]=my_plane_rank + i*ranks_per_plane;
        }

        // Create the new inter-plane group
        MPI_Group intpl_group;
        MPI_Group_incl(group, nplanes, intpl_ranks.data(), &intpl_group);

        // Create the new inter-plane communicator
        MPI_Comm_create(comm, intpl_group, &intpl_comm);

        MPI_Comm_rank(intpl_comm, &my_intpl_rank);
        MPI_Comm_size(intpl_comm, &n_intpl_ranks );

        /*
        call check_point('adios mpi communication?')
        do i=0, sml_intpl_nranks-1
        adios_ranks(i)= i*ranks_per_plane
        enddo
        call mpi_group_incl(sml_comm_group, sml_intpl_nranks, adios_ranks, sml_adios_group, ierr)
        call mpi_comm_create(sml_comm, sml_adios_group, sml_adios_comm, ierr)


    //#ifndef NO_TASKMAP
        // output task-to-node mapping
        write(c_color,'(i8)') sml_comm_color
        if (sml_mype .eq. 0) then
        open(unit=14, file='TASKMAP_Color'//trim(adjustl(c_color))//'.txt', &
        status='OLD', access='SEQUENTIAL', position='APPEND' )
        endif
        call taskmap_write(14, sml_comm, &
        'SML_COMM COLOR #'//trim(adjustl(c_color)), sml_plane_mype, sml_intpl_mype, .false.)
        if (sml_mype .eq. 0) close(14)
    //#endif
    //#ifdef XGC1
        //jyc print mpi placement info
        if (sml_verbose) then
        call mpi_get_processor_name(nodename, len, ierr)
        print *, sml_mype, 'Process (plane,node,color):', sml_plane_index, trim(nodename), ' ', sml_comm_color
        endif
    //#endif
        */
    }
};

// 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