mpi_accumulate.h

Go to the documentation of this file.

 
#pragma once
#include <mpi.h>
#include <array>
#include <vector>
#include <map>
#include "accumulate.h"
 
namespace dg
{
namespace exblas {
 
inline void mpi_reduce_communicator(MPI_Comm comm,
    MPI_Comm* comm_mod, MPI_Comm* comm_mod_reduce){
    //we keep track of communicators that were created in the past
    // (static function variables are constructed the program
    // encounters them and destruced at program termination)
    static std::map<MPI_Comm, std::array<MPI_Comm, 2>> comm_mods;
    assert( comm != MPI_COMM_NULL);
    if( comm_mods.count(comm) == 1 )
    {
        *comm_mod = comm_mods[comm][0];
        *comm_mod_reduce = comm_mods[comm][1];
        return;
    }
    int rank, size;
    MPI_Comm_rank( comm, &rank);
    MPI_Comm_size( comm, &size);
    // For example
    // xxxxxxxxxxxxxxxx // 16 (comm_mod)
    // xxxxxxxxxxxxxxxx // 16
    // xxxxxx           //  6
    // 3333332222222222 // (comm_mod_reduce)
    int mod = 128;
    // 0-127, 128-255, ...
    MPI_Comm_split( comm, rank/mod, rank%mod, comm_mod); //collective call
    // Here we split rank%mod instead of returning MPI_COMM_NULL for != 0
    // https://github.com/open-mpi/ompi/issues/13081
    MPI_Comm_split( comm, rank%mod, rank, comm_mod_reduce);
 
    comm_mods[comm] = {*comm_mod, *comm_mod_reduce};
}
 
inline void reduce_mpi_cpu(  unsigned num_superacc, int64_t* in, int64_t* out,
MPI_Comm comm, MPI_Comm comm_mod, MPI_Comm comm_mod_reduce )
{
    for( unsigned i=0; i<num_superacc; i++)
    {
        int imin=exblas::IMIN, imax=exblas::IMAX;
        cpu::Normalize(&in[i*exblas::BIN_COUNT], imin, imax);
    }
    int size_comm, size_comm_mod;
    MPI_Comm_size( comm, &size_comm);
    MPI_Comm_size( comm_mod, &size_comm_mod);
    if( size_comm == size_comm_mod) // size < 128
    {
        MPI_Allreduce( in, out, num_superacc*exblas::BIN_COUNT, MPI_LONG, MPI_SUM, comm);
        return;
    }
 
    MPI_Reduce(in, out, num_superacc*exblas::BIN_COUNT, MPI_LONG, MPI_SUM, 0,
        comm_mod);
    int rank;
    MPI_Comm_rank( comm_mod, &rank);
    if(rank == 0)
    {
        int size;
        MPI_Comm_size( comm_mod_reduce, &size);
        if( size > 1)
        {
            for( unsigned i=0; i<num_superacc; i++)
            {
                int imin=exblas::IMIN, imax=exblas::IMAX;
                cpu::Normalize(&out[i*exblas::BIN_COUNT], imin, imax);
                for( int k=0; k<exblas::BIN_COUNT; k++)
                    in[i*BIN_COUNT+k] = out[i*BIN_COUNT+k];
            }
            MPI_Reduce(in, out, num_superacc*exblas::BIN_COUNT, MPI_LONG,
                MPI_SUM, 0, comm_mod_reduce);
        }
    }
    MPI_Bcast( out, num_superacc*exblas::BIN_COUNT, MPI_LONG, 0, comm);
}
 
}//namespace exblas
} //namespace dg