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

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

EP update all

File size: 5.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
13using namespace std;
14
15namespace ep_lib
16{
17
18  int MPI_Gatherv_local(const void *sendbuf, int count, MPI_Datatype datatype, void *recvbuf, const int recvcounts[], const int displs[], int local_root, MPI_Comm comm)
19  {
20    assert(valid_type(datatype));
21
22    ::MPI_Aint datasize, lb;
23    ::MPI_Type_get_extent(to_mpi_type(datatype), &lb, &datasize);
24
25    int ep_rank_loc = comm.ep_comm_ptr->size_rank_info[1].first;
26    int num_ep = comm.ep_comm_ptr->size_rank_info[1].second;
27
28    //if(ep_rank_loc == local_root) printf("local_gatherv : recvcounts = %d %d\n\n", recvcounts[0], recvcounts[1]);
29    //if(ep_rank_loc == local_root) printf("local_gatherv : displs = %d %d\n\n", displs[0], displs[1]);
30
31    #pragma omp critical (_gatherv)
32    comm.my_buffer->void_buffer[ep_rank_loc] = const_cast< void* >(sendbuf);
33
34    MPI_Barrier_local(comm);
35
36    if(ep_rank_loc == local_root)
37    {
38      for(int i=0; i<num_ep; i++)
39        memcpy(recvbuf + datasize*displs[i], comm.my_buffer->void_buffer[i], datasize*recvcounts[i]);
40
41    }
42
43    MPI_Barrier_local(comm);
44  }
45
46  int MPI_Gatherv(const void *sendbuf, int sendcount, MPI_Datatype sendtype, void *recvbuf, const int input_recvcounts[], const int input_displs[],
47                  MPI_Datatype recvtype, int root, MPI_Comm comm)
48  {
49 
50    if(!comm.is_ep)
51    {
52      ::MPI_Gatherv(const_cast<void*>(sendbuf), sendcount, static_cast< ::MPI_Datatype>(sendtype), recvbuf, const_cast<int*>(input_recvcounts), const_cast<int*>(input_displs),
53                    static_cast< ::MPI_Datatype>(recvtype), root, static_cast< ::MPI_Comm>(comm.mpi_comm));
54      return 0;
55    }
56
57
58    assert(sendtype == recvtype);
59
60   
61    int ep_rank = comm.ep_comm_ptr->size_rank_info[0].first;
62    int ep_rank_loc = comm.ep_comm_ptr->size_rank_info[1].first;
63    int mpi_rank = comm.ep_comm_ptr->size_rank_info[2].first;
64    int ep_size = comm.ep_comm_ptr->size_rank_info[0].second;
65    int num_ep = comm.ep_comm_ptr->size_rank_info[1].second;
66    int mpi_size = comm.ep_comm_ptr->size_rank_info[2].second;
67
68    int root_mpi_rank = comm.rank_map->at(root).second;
69    int root_ep_loc = comm.rank_map->at(root).first;
70
71    ::MPI_Aint datasize, lb;
72    ::MPI_Type_get_extent(to_mpi_type(sendtype), &lb, &datasize);
73
74    int *recvcounts;
75    int* displs;
76
77    recvcounts = new int[ep_size];
78    displs = new int[ep_size];
79
80
81    bool is_master = (ep_rank_loc==0 && mpi_rank != root_mpi_rank ) || ep_rank == root;
82    bool is_root = ep_rank == root;
83
84    void* local_recvbuf;
85    std::vector<int>local_recvcounts(num_ep, 0);
86    std::vector<int>local_displs(num_ep, 0);
87
88
89    if(is_root)
90    { 
91      copy(input_recvcounts, input_recvcounts+ep_size, recvcounts);
92      copy(input_displs, input_displs+ep_size, displs);
93    }
94
95    MPI_Bcast(recvcounts, ep_size, MPI_INT, root, comm);
96    MPI_Bcast(displs, ep_size, MPI_INT, root, comm);
97
98    if(mpi_rank == root_mpi_rank) MPI_Gather_local(&sendcount, 1, MPI_INT, local_recvcounts.data(), root_ep_loc, comm);
99    else                          MPI_Gather_local(&sendcount, 1, MPI_INT, local_recvcounts.data(), 0, comm);
100
101
102
103    if(is_master)
104    {
105      int local_recvbuf_size = std::accumulate(local_recvcounts.begin(), local_recvcounts.end(), 0);
106     
107      for(int i=1; i<num_ep; i++)
108        local_displs[i] = local_displs[i-1] + local_recvcounts[i-1];
109
110      local_recvbuf = new void*[datasize * local_recvbuf_size];
111    }
112
113    if(mpi_rank == root_mpi_rank) MPI_Gatherv_local(sendbuf, sendcount, sendtype, local_recvbuf, local_recvcounts.data(), local_displs.data(), root_ep_loc, comm);
114    else                          MPI_Gatherv_local(sendbuf, sendcount, sendtype, local_recvbuf, local_recvcounts.data(), local_displs.data(), 0, comm);
115
116    //if(is_master) printf("local_recvbuf = %d %d %d %d\n", static_cast<int*>(local_recvbuf)[0], static_cast<int*>(local_recvbuf)[1], static_cast<int*>(local_recvbuf)[2], static_cast<int*>(local_recvbuf)[3]);
117
118    void* tmp_recvbuf;
119    int tmp_recvbuf_size = std::accumulate(recvcounts, recvcounts+ep_size, 0);
120
121    if(is_root) tmp_recvbuf = new void*[datasize * tmp_recvbuf_size];
122
123
124    std::vector<int> mpi_recvcounts(mpi_size, 0);
125    std::vector<int> mpi_displs(mpi_size, 0);
126
127
128    if(is_master)
129    {
130      for(int i=0; i<ep_size; i++)
131      {
132        mpi_recvcounts[comm.rank_map->at(i).second]+=recvcounts[i];
133      }
134
135      for(int i=1; i<mpi_size; i++)
136        mpi_displs[i] = mpi_displs[i-1] + mpi_recvcounts[i-1];
137
138
139      ::MPI_Gatherv(local_recvbuf, sendcount*num_ep, sendtype, tmp_recvbuf, mpi_recvcounts.data(), mpi_displs.data(), recvtype, root_mpi_rank, to_mpi_comm(comm.mpi_comm));
140    }   
141
142
143    // reorder data
144    if(is_root)
145    {
146      int offset;
147      for(int i=0; i<ep_size; i++)
148      {
149        int extra = 0;
150        for(int j=0, k=0; j<ep_size, k<comm.rank_map->at(i).first; j++)
151          if(comm.rank_map->at(i).second == comm.rank_map->at(j).second)
152          {
153            extra += recvcounts[j];
154            k++;
155          } 
156
157        offset = mpi_displs[comm.rank_map->at(i).second] +  extra;
158
159        memcpy(recvbuf+displs[i]*datasize, tmp_recvbuf+offset*datasize, recvcounts[i]*datasize);
160       
161      }
162
163    }
164
165    delete[] recvcounts;
166    delete[] displs;
167
168    if(is_master)
169    {
170      delete[] local_recvbuf;
171    }
172    if(is_root) delete[] tmp_recvbuf;
173  }
174
175}
Note: See TracBrowser for help on using the repository browser.