xref: /petsc/src/sys/utils/mpimesg.c (revision 9895aa37ac365bac650f6bd8bf977519f7222510)
1 
2 #include <petscsys.h>        /*I  "petscsys.h"  I*/
3 
4 
5 #undef __FUNCT__
6 #define __FUNCT__ "PetscGatherNumberOfMessages"
7 /*@C
8   PetscGatherNumberOfMessages -  Computes the number of messages a node expects to receive
9 
10   Collective on MPI_Comm
11 
12   Input Parameters:
13 + comm     - Communicator
14 . iflags   - an array of integers of length sizeof(comm). A '1' in ilengths[i] represent a
15              message from current node to ith node. Optionally NULL
16 - ilengths - Non zero ilengths[i] represent a message to i of length ilengths[i].
17              Optionally NULL.
18 
19   Output Parameters:
20 . nrecvs    - number of messages received
21 
22   Level: developer
23 
24   Concepts: mpi utility
25 
26   Notes:
27   With this info, the correct message lengths can be determined using
28   PetscGatherMessageLengths()
29 
30   Either iflags or ilengths should be provided.  If iflags is not
31   provided (NULL) it can be computed from ilengths. If iflags is
32   provided, ilengths is not required.
33 
34 .seealso: PetscGatherMessageLengths()
35 @*/
36 PetscErrorCode  PetscGatherNumberOfMessages(MPI_Comm comm,const PetscMPIInt iflags[],const PetscMPIInt ilengths[],PetscMPIInt *nrecvs)
37 {
38   PetscMPIInt    size,rank,*recv_buf,i,*iflags_local = NULL,*iflags_localm = NULL;
39   PetscErrorCode ierr;
40 
41   PetscFunctionBegin;
42   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
43   ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
44 
45   ierr = PetscMalloc2(size,PetscMPIInt,&recv_buf,size,PetscMPIInt,&iflags_localm);CHKERRQ(ierr);
46 
47   /* If iflags not provided, compute iflags from ilengths */
48   if (!iflags) {
49     if (!ilengths) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Either iflags or ilengths should be provided");
50     iflags_local = iflags_localm;
51     for (i=0; i<size; i++) {
52       if (ilengths[i]) iflags_local[i] = 1;
53       else iflags_local[i] = 0;
54     }
55   } else iflags_local = (PetscMPIInt*) iflags;
56 
57   /* Post an allreduce to determine the numer of messages the current node will receive */
58   ierr    = MPI_Allreduce(iflags_local,recv_buf,size,MPI_INT,MPI_SUM,comm);CHKERRQ(ierr);
59   *nrecvs = recv_buf[rank];
60 
61   ierr = PetscFree2(recv_buf,iflags_localm);CHKERRQ(ierr);
62   PetscFunctionReturn(0);
63 }
64 
65 
66 #undef __FUNCT__
67 #define __FUNCT__ "PetscGatherMessageLengths"
68 /*@C
69   PetscGatherMessageLengths - Computes info about messages that a MPI-node will receive,
70   including (from-id,length) pairs for each message.
71 
72   Collective on MPI_Comm
73 
74   Input Parameters:
75 + comm      - Communicator
76 . nsends    - number of messages that are to be sent.
77 . nrecvs    - number of messages being received
78 - ilengths  - an array of integers of length sizeof(comm)
79               a non zero ilengths[i] represent a message to i of length ilengths[i]
80 
81 
82   Output Parameters:
83 + onodes    - list of node-ids from which messages are expected
84 - olengths  - corresponding message lengths
85 
86   Level: developer
87 
88   Concepts: mpi utility
89 
90   Notes:
91   With this info, the correct MPI_Irecv() can be posted with the correct
92   from-id, with a buffer with the right amount of memory required.
93 
94   The calling function deallocates the memory in onodes and olengths
95 
96   To determine nrecevs, one can use PetscGatherNumberOfMessages()
97 
98 .seealso: PetscGatherNumberOfMessages()
99 @*/
100 PetscErrorCode  PetscGatherMessageLengths(MPI_Comm comm,PetscMPIInt nsends,PetscMPIInt nrecvs,const PetscMPIInt ilengths[],PetscMPIInt **onodes,PetscMPIInt **olengths)
101 {
102   PetscErrorCode ierr;
103   PetscMPIInt    size,tag,i,j;
104   MPI_Request    *s_waits  = NULL,*r_waits = NULL;
105   MPI_Status     *w_status = NULL;
106 
107   PetscFunctionBegin;
108   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
109   ierr = PetscCommGetNewTag(comm,&tag);CHKERRQ(ierr);
110 
111   /* cannot use PetscMalloc3() here because in the call to MPI_Waitall() they MUST be contiguous */
112   ierr    = PetscMalloc2(nrecvs+nsends,MPI_Request,&r_waits,nrecvs+nsends,MPI_Status,&w_status);CHKERRQ(ierr);
113   s_waits = r_waits+nrecvs;
114 
115   /* Post the Irecv to get the message length-info */
116   ierr = PetscMalloc(nrecvs*sizeof(PetscMPIInt),olengths);CHKERRQ(ierr);
117   for (i=0; i<nrecvs; i++) {
118     ierr = MPI_Irecv((*olengths)+i,1,MPI_INT,MPI_ANY_SOURCE,tag,comm,r_waits+i);CHKERRQ(ierr);
119   }
120 
121   /* Post the Isends with the message length-info */
122   for (i=0,j=0; i<size; ++i) {
123     if (ilengths[i]) {
124       ierr = MPI_Isend((void*)(ilengths+i),1,MPI_INT,i,tag,comm,s_waits+j);CHKERRQ(ierr);
125       j++;
126     }
127   }
128 
129   /* Post waits on sends and receivs */
130   if (nrecvs+nsends) {ierr = MPI_Waitall(nrecvs+nsends,r_waits,w_status);CHKERRQ(ierr);}
131 
132   /* Pack up the received data */
133   ierr = PetscMalloc(nrecvs*sizeof(PetscMPIInt),onodes);CHKERRQ(ierr);
134   for (i=0; i<nrecvs; ++i) (*onodes)[i] = w_status[i].MPI_SOURCE;
135   ierr = PetscFree2(r_waits,w_status);CHKERRQ(ierr);
136   PetscFunctionReturn(0);
137 }
138 
139 #undef __FUNCT__
140 #define __FUNCT__ "PetscGatherMessageLengths2"
141 /*@C
142   PetscGatherMessageLengths2 - Computes info about messages that a MPI-node will receive,
143   including (from-id,length) pairs for each message. Same functionality as PetscGatherMessageLengths()
144   except it takes TWO ilenths and output TWO olengths.
145 
146   Collective on MPI_Comm
147 
148   Input Parameters:
149 + comm      - Communicator
150 . nsends    - number of messages that are to be sent.
151 . nrecvs    - number of messages being received
152 - ilengths1, ilengths2 - array of integers of length sizeof(comm)
153               a non zero ilengths[i] represent a message to i of length ilengths[i]
154 
155   Output Parameters:
156 + onodes    - list of node-ids from which messages are expected
157 - olengths1, olengths2 - corresponding message lengths
158 
159   Level: developer
160 
161   Concepts: mpi utility
162 
163   Notes:
164   With this info, the correct MPI_Irecv() can be posted with the correct
165   from-id, with a buffer with the right amount of memory required.
166 
167   The calling function deallocates the memory in onodes and olengths
168 
169   To determine nrecevs, one can use PetscGatherNumberOfMessages()
170 
171 .seealso: PetscGatherMessageLengths() and PetscGatherNumberOfMessages()
172 @*/
173 PetscErrorCode  PetscGatherMessageLengths2(MPI_Comm comm,PetscMPIInt nsends,PetscMPIInt nrecvs,const PetscMPIInt ilengths1[],const PetscMPIInt ilengths2[],PetscMPIInt **onodes,PetscMPIInt **olengths1,PetscMPIInt **olengths2)
174 {
175   PetscErrorCode ierr;
176   PetscMPIInt    size,tag,i,j,*buf_s = NULL,*buf_r = NULL,*buf_j = NULL;
177   MPI_Request    *s_waits  = NULL,*r_waits = NULL;
178   MPI_Status     *w_status = NULL;
179 
180   PetscFunctionBegin;
181   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
182   ierr = PetscCommGetNewTag(comm,&tag);CHKERRQ(ierr);
183 
184   /* cannot use PetscMalloc5() because r_waits and s_waits must be contiquous for the call to MPI_Waitall() */
185   ierr = PetscMalloc4(nrecvs+nsends,MPI_Request,&r_waits,2*nrecvs,PetscMPIInt,&buf_r,2*nsends,PetscMPIInt,&buf_s,nrecvs+nsends,MPI_Status,&w_status);CHKERRQ(ierr);
186   s_waits = r_waits + nrecvs;
187 
188   /* Post the Irecv to get the message length-info */
189   ierr = PetscMalloc((nrecvs+1)*sizeof(PetscMPIInt),olengths1);CHKERRQ(ierr);
190   ierr = PetscMalloc((nrecvs+1)*sizeof(PetscMPIInt),olengths2);CHKERRQ(ierr);
191   for (i=0; i<nrecvs; i++) {
192     buf_j = buf_r + (2*i);
193     ierr  = MPI_Irecv(buf_j,2,MPI_INT,MPI_ANY_SOURCE,tag,comm,r_waits+i);CHKERRQ(ierr);
194   }
195 
196   /* Post the Isends with the message length-info */
197   for (i=0,j=0; i<size; ++i) {
198     if (ilengths1[i]) {
199       buf_j    = buf_s + (2*j);
200       buf_j[0] = *(ilengths1+i);
201       buf_j[1] = *(ilengths2+i);
202       ierr = MPI_Isend(buf_j,2,MPI_INT,i,tag,comm,s_waits+j);CHKERRQ(ierr);
203       j++;
204     }
205   }
206 
207   /* Post waits on sends and receivs */
208   if (nrecvs+nsends) {ierr = MPI_Waitall(nrecvs+nsends,r_waits,w_status);CHKERRQ(ierr);}
209 
210 
211   /* Pack up the received data */
212   ierr = PetscMalloc((nrecvs+1)*sizeof(PetscMPIInt),onodes);CHKERRQ(ierr);
213   for (i=0; i<nrecvs; ++i) {
214     (*onodes)[i]    = w_status[i].MPI_SOURCE;
215     buf_j           = buf_r + (2*i);
216     (*olengths1)[i] = buf_j[0];
217     (*olengths2)[i] = buf_j[1];
218   }
219 
220   ierr = PetscFree4(r_waits,buf_r,buf_s,w_status);CHKERRQ(ierr);
221   PetscFunctionReturn(0);
222 }
223 
224 /*
225 
226   Allocate a bufffer sufficient to hold messages of size specified in olengths.
227   And post Irecvs on these buffers using node info from onodes
228 
229  */
230 #undef __FUNCT__
231 #define __FUNCT__ "PetscPostIrecvInt"
232 PetscErrorCode  PetscPostIrecvInt(MPI_Comm comm,PetscMPIInt tag,PetscMPIInt nrecvs,const PetscMPIInt onodes[],const PetscMPIInt olengths[],PetscInt ***rbuf,MPI_Request **r_waits)
233 {
234   PetscErrorCode ierr;
235   PetscInt       **rbuf_t,i,len = 0;
236   MPI_Request    *r_waits_t;
237 
238   PetscFunctionBegin;
239   /* compute memory required for recv buffers */
240   for (i=0; i<nrecvs; i++) len += olengths[i];  /* each message length */
241 
242   /* allocate memory for recv buffers */
243   ierr = PetscMalloc((nrecvs+1)*sizeof(PetscInt*),&rbuf_t);CHKERRQ(ierr);
244   ierr = PetscMalloc(len*sizeof(PetscInt),&rbuf_t[0]);CHKERRQ(ierr);
245   for (i=1; i<nrecvs; ++i) rbuf_t[i] = rbuf_t[i-1] + olengths[i-1];
246 
247   /* Post the receives */
248   ierr = PetscMalloc(nrecvs*sizeof(MPI_Request),&r_waits_t);CHKERRQ(ierr);
249   for (i=0; i<nrecvs; ++i) {
250     ierr = MPI_Irecv(rbuf_t[i],olengths[i],MPIU_INT,onodes[i],tag,comm,r_waits_t+i);CHKERRQ(ierr);
251   }
252 
253   *rbuf    = rbuf_t;
254   *r_waits = r_waits_t;
255   PetscFunctionReturn(0);
256 }
257 
258 #undef __FUNCT__
259 #define __FUNCT__ "PetscPostIrecvScalar"
260 PetscErrorCode  PetscPostIrecvScalar(MPI_Comm comm,PetscMPIInt tag,PetscMPIInt nrecvs,const PetscMPIInt onodes[],const PetscMPIInt olengths[],PetscScalar ***rbuf,MPI_Request **r_waits)
261 {
262   PetscErrorCode ierr;
263   PetscMPIInt    i;
264   PetscScalar    **rbuf_t;
265   MPI_Request    *r_waits_t;
266   PetscInt       len = 0;
267 
268   PetscFunctionBegin;
269   /* compute memory required for recv buffers */
270   for (i=0; i<nrecvs; i++) len += olengths[i];  /* each message length */
271 
272   /* allocate memory for recv buffers */
273   ierr = PetscMalloc((nrecvs+1)*sizeof(PetscScalar*),&rbuf_t);CHKERRQ(ierr);
274   ierr = PetscMalloc(len*sizeof(PetscScalar),&rbuf_t[0]);CHKERRQ(ierr);
275   for (i=1; i<nrecvs; ++i) rbuf_t[i] = rbuf_t[i-1] + olengths[i-1];
276 
277   /* Post the receives */
278   ierr = PetscMalloc(nrecvs*sizeof(MPI_Request),&r_waits_t);CHKERRQ(ierr);
279   for (i=0; i<nrecvs; ++i) {
280     ierr = MPI_Irecv(rbuf_t[i],olengths[i],MPIU_SCALAR,onodes[i],tag,comm,r_waits_t+i);CHKERRQ(ierr);
281   }
282 
283   *rbuf    = rbuf_t;
284   *r_waits = r_waits_t;
285   PetscFunctionReturn(0);
286 }
287