source: XIOS/dev/branch_openmp/extern/src_ep_dev/ep_allgatherv.cpp @ 1295

Last change on this file since 1295 was 1295, checked in by yushan, 4 years ago

EP update all

File size: 3.4 KB
Line 
1/*!
2   \file ep_gather.cpp
3   \since 2 may 2016
4
5   \brief Definitions of MPI collective function: MPI_Gatherv, MPI_Allgatherv
6 */
7
8#include "ep_lib.hpp"
9#include <mpi.h>
10#include "ep_declaration.hpp"
11#include "ep_mpi.hpp"
12
13
14using namespace std;
15
16namespace ep_lib
17{
18
19  int MPI_Allgatherv(const void *sendbuf, int sendcount, MPI_Datatype sendtype, void *recvbuf, const int recvcounts[], const int displs[], MPI_Datatype recvtype, MPI_Comm comm)
20  {
21
22    if(!comm.is_ep && comm.mpi_comm)
23    {
24      ::MPI_Allgatherv(sendbuf, sendcount, to_mpi_type(sendtype), recvbuf, recvcounts, displs, to_mpi_type(recvtype), to_mpi_comm(comm.mpi_comm));
25      return 0;
26    }
27
28    if(!comm.mpi_comm) return 0;
29
30
31    assert(valid_type(sendtype) && valid_type(recvtype));
32
33    MPI_Datatype datatype = sendtype;
34    int count = sendcount;
35
36    ::MPI_Aint datasize, lb;
37
38    ::MPI_Type_get_extent(to_mpi_type(datatype), &lb, &datasize);
39
40
41    int ep_rank = comm.ep_comm_ptr->size_rank_info[0].first;
42    int ep_rank_loc = comm.ep_comm_ptr->size_rank_info[1].first;
43    int mpi_rank = comm.ep_comm_ptr->size_rank_info[2].first;
44    int ep_size = comm.ep_comm_ptr->size_rank_info[0].second;
45    int num_ep = comm.ep_comm_ptr->size_rank_info[1].second;
46    int mpi_size = comm.ep_comm_ptr->size_rank_info[2].second;
47
48    assert(sendcount == recvcounts[ep_rank]);
49
50    bool is_master = ep_rank_loc==0;
51
52    void* local_recvbuf;
53    void* tmp_recvbuf;
54
55    int recvbuf_size = 0;
56    for(int i=0; i<ep_size; i++)
57      recvbuf_size = max(recvbuf_size, displs[i]+recvcounts[i]);
58
59
60    vector<int>local_recvcounts(num_ep, 0);
61    vector<int>local_displs(num_ep, 0);
62
63    MPI_Gather_local(&sendcount, 1, MPI_INT, local_recvcounts.data(), 0, comm);
64    for(int i=1; i<num_ep; i++) local_displs[i] = local_displs[i-1] + local_recvcounts[i-1]; 
65
66
67    if(is_master)
68    {
69      local_recvbuf = new void*[datasize * std::accumulate(local_recvcounts.begin(), local_recvcounts.end(), 0)];
70      tmp_recvbuf = new void*[datasize * std::accumulate(recvcounts, recvcounts+ep_size, 0)];
71    }
72
73    MPI_Gatherv_local(sendbuf, count, datatype, local_recvbuf, local_recvcounts.data(), local_displs.data(), 0, comm);
74
75
76    if(is_master)
77    {
78      std::vector<int>mpi_recvcounts(mpi_size, 0);
79      std::vector<int>mpi_displs(mpi_size, 0);
80
81      int local_sendcount = std::accumulate(local_recvcounts.begin(), local_recvcounts.end(), 0);
82      MPI_Allgather(&local_sendcount, 1, MPI_INT, mpi_recvcounts.data(), 1, MPI_INT, to_mpi_comm(comm.mpi_comm));
83
84      for(int i=1; i<mpi_size; i++)
85        mpi_displs[i] = mpi_displs[i-1] + mpi_recvcounts[i-1];
86
87
88      ::MPI_Allgatherv(local_recvbuf, local_sendcount, to_mpi_type(datatype), tmp_recvbuf, mpi_recvcounts.data(), mpi_displs.data(), to_mpi_type(datatype), to_mpi_comm(comm.mpi_comm));
89
90      // reorder
91      int offset;
92      for(int i=0; i<ep_size; i++)
93      {
94        int extra = 0;
95        for(int j=0, k=0; j<ep_size, k<comm.rank_map->at(i).first; j++)
96          if(comm.rank_map->at(i).second == comm.rank_map->at(j).second)
97          {
98            extra += recvcounts[j];
99            k++;
100          } 
101
102        offset = mpi_displs[comm.rank_map->at(i).second] +  extra;
103
104        memcpy(recvbuf+displs[i]*datasize, tmp_recvbuf+offset*datasize, recvcounts[i]*datasize);
105       
106      }
107
108    }
109
110    MPI_Bcast_local(recvbuf, recvbuf_size, datatype, 0, comm);
111
112    if(is_master)
113    {
114      delete[] local_recvbuf;
115      delete[] tmp_recvbuf;
116    }
117
118
119  }
120
121
122
123}
Note: See TracBrowser for help on using the repository browser.