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