source: XIOS/dev/branch_openmp/extern/ep_dev/ep_scatterv.cpp @ 1503

Last change on this file since 1503 was 1503, checked in by yushan, 3 years ago

rank_map is passed from vector to map, in order to have more flexibility in comm_split

File size: 5.8 KB
Line 
1/*!
2   \file ep_gather.cpp
3   \since 2 may 2016
4
5   \brief Definitions of MPI collective function: MPI_Scatterv
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_Scatterv_local(const void *sendbuf, const int sendcounts[], const int displs[], MPI_Datatype sendtype, void *recvbuf, int recvcount,
19                   MPI_Datatype recvtype, int local_root, MPI_Comm comm)
20  {
21
22    assert(valid_type(sendtype) && valid_type(recvtype));
23
24    ::MPI_Aint datasize, lb;
25    ::MPI_Type_get_extent(to_mpi_type(sendtype), &lb, &datasize);
26
27    int ep_rank_loc = comm->ep_comm_ptr->size_rank_info[1].first;
28    int num_ep = comm->ep_comm_ptr->size_rank_info[1].second;
29
30    assert(recvcount == sendcounts[ep_rank_loc]);
31
32    if(ep_rank_loc == local_root)
33      comm->my_buffer->void_buffer[local_root] = const_cast<void*>(sendbuf);
34
35    MPI_Barrier_local(comm);
36
37    #pragma omp critical (_scatterv)     
38    memcpy(recvbuf, comm->my_buffer->void_buffer[local_root]+datasize*displs[ep_rank_loc], datasize * recvcount);
39   
40
41    MPI_Barrier_local(comm);
42  }
43
44  int MPI_Scatterv(const void *sendbuf, const int sendcounts[], const int displs[], MPI_Datatype sendtype, void *recvbuf, int recvcount,
45                   MPI_Datatype recvtype, int root, MPI_Comm comm)
46  {
47    if(!comm->is_ep)
48    {
49      return ::MPI_Scatterv(sendbuf, sendcounts, displs, to_mpi_type(sendtype), recvbuf, recvcount, to_mpi_type(recvtype), root, to_mpi_comm(comm->mpi_comm));
50    }
51   
52    assert(sendtype == recvtype);
53
54    int ep_rank = comm->ep_comm_ptr->size_rank_info[0].first;
55    int ep_rank_loc = comm->ep_comm_ptr->size_rank_info[1].first;
56    int mpi_rank = comm->ep_comm_ptr->size_rank_info[2].first;
57    int ep_size = comm->ep_comm_ptr->size_rank_info[0].second;
58    int num_ep = comm->ep_comm_ptr->size_rank_info[1].second;
59    int mpi_size = comm->ep_comm_ptr->size_rank_info[2].second;
60
61    int root_mpi_rank = comm->ep_rank_map->at(root).second;
62    int root_ep_loc = comm->ep_rank_map->at(root).first;
63
64    bool is_master = (ep_rank_loc==0 && mpi_rank != root_mpi_rank ) || ep_rank == root;
65    bool is_root = ep_rank == root;
66
67    MPI_Datatype datatype = sendtype;
68    int count = recvcount;
69
70    ::MPI_Aint datasize, lb;
71    ::MPI_Type_get_extent(to_mpi_type(datatype), &lb, &datasize);
72   
73    void *tmp_sendbuf;
74    if(is_root) tmp_sendbuf = new void*[ep_size * count * datasize];
75
76    // reorder tmp_sendbuf
77    std::vector<int>local_ranks(num_ep);
78    std::vector<int>ranks(ep_size);
79
80    if(mpi_rank == root_mpi_rank) MPI_Gather_local(&ep_rank, 1, MPI_INT, local_ranks.data(), root_ep_loc, comm);
81    else                          MPI_Gather_local(&ep_rank, 1, MPI_INT, local_ranks.data(), 0, comm);
82
83
84    std::vector<int> recvcounts(mpi_size, 0);
85    std::vector<int> my_displs(mpi_size, 0);
86
87
88    if(is_master)
89    {
90      for(int i=0; i<ep_size; i++)
91      {
92        recvcounts[comm->ep_rank_map->at(i).second]++;
93      }
94
95      for(int i=1; i<mpi_size; i++)
96        my_displs[i] = my_displs[i-1] + recvcounts[i-1];
97
98      ::MPI_Gatherv(local_ranks.data(), num_ep, to_mpi_type(MPI_INT), ranks.data(), recvcounts.data(), my_displs.data(), to_mpi_type(MPI_INT), root_mpi_rank, to_mpi_comm(comm->mpi_comm));
99    }
100
101
102    if(is_root)
103    {
104      int local_displs = 0;
105      for(int i=0; i<ep_size; i++)
106      {
107        //printf("i=%d : start from %d, src displs = %d, count = %d\n ", i, local_displs/datasize, displs[ranks[i]], sendcounts[ranks[i]]);
108        memcpy(tmp_sendbuf+local_displs, sendbuf + displs[ranks[i]]*datasize, sendcounts[ranks[i]]*datasize);
109        local_displs += sendcounts[ranks[i]]*datasize;
110      }
111     
112      //for(int i=0; i<ep_size*2; i++) printf("%d\t", static_cast<int*>(const_cast<void*>(tmp_sendbuf))[i]);
113    }
114
115    // MPI_Scatterv from root to masters
116 
117    void* local_sendbuf;
118    int local_sendcount;
119    if(mpi_rank == root_mpi_rank) MPI_Reduce_local(&recvcount, &local_sendcount, 1, MPI_INT, MPI_SUM, root_ep_loc, comm);
120    else                          MPI_Reduce_local(&recvcount, &local_sendcount, 1, MPI_INT, MPI_SUM, 0, comm);
121
122    if(is_master) 
123    {
124      local_sendbuf = new void*[datasize * local_sendcount];
125 
126      ::MPI_Gather(&local_sendcount, 1, to_mpi_type(MPI_INT), recvcounts.data(), 1, to_mpi_type(MPI_INT), root_mpi_rank, to_mpi_comm(comm->mpi_comm));
127
128      if(is_root) for(int i=1; i<mpi_size; i++) my_displs[i] = my_displs[i-1] + recvcounts[i-1];
129
130      ::MPI_Scatterv(tmp_sendbuf, recvcounts.data(), my_displs.data(), to_mpi_type(sendtype), local_sendbuf, num_ep*count, to_mpi_type(recvtype), root_mpi_rank, to_mpi_comm(comm->mpi_comm));
131
132      // printf("my_displs = %d %d %d %d\n", my_displs[0], my_displs[1], my_displs[2], my_displs[3] );
133
134      // printf("%d %d %d %d %d %d %d %d\n", static_cast<int*>(local_sendbuf)[0], static_cast<int*>(local_sendbuf)[1], static_cast<int*>(local_sendbuf)[2], static_cast<int*>(local_sendbuf)[3],
135      //                                     static_cast<int*>(local_sendbuf)[4], static_cast<int*>(local_sendbuf)[5], static_cast<int*>(local_sendbuf)[6], static_cast<int*>(local_sendbuf)[7]);
136    }
137
138    std::vector<int>local_sendcounts(num_ep, 0);
139    std::vector<int>local_displs(num_ep, 0);
140
141    MPI_Gather_local(&recvcount, 1, MPI_INT, local_sendcounts.data(), 0, comm);
142    MPI_Bcast_local(local_sendcounts.data(), num_ep, MPI_INT, 0, comm);
143    for(int i=1; i<num_ep; i++)
144      local_displs[i] = local_displs[i-1] + local_sendcounts[i-1];
145   
146
147    if(mpi_rank == root_mpi_rank) MPI_Scatterv_local(local_sendbuf, local_sendcounts.data(), local_displs.data(), sendtype, recvbuf, recvcount, recvtype, root_ep_loc, comm);
148    else                          MPI_Scatterv_local(local_sendbuf, local_sendcounts.data(), local_displs.data(), sendtype, recvbuf, recvcount, recvtype, 0, comm);
149
150
151    if(is_root) delete[] tmp_sendbuf;
152    if(is_master) delete[] local_sendbuf;
153  }
154
155
156}
Note: See TracBrowser for help on using the repository browser.