xref: /petsc/src/sys/utils/mpimesg.c (revision 9596e0b48258fba4fca4f68feb5185896facfe69)
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,const PetscMPIInt iflags[],const 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 = (PetscMPIInt *) 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,const 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((void*)(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 
145 #undef __FUNCT__
146 #define __FUNCT__ "PetscGatherMessageLengths2"
147 /*@C
148   PetscGatherMessageLengths2 - Computes info about messages that a MPI-node will receive,
149   including (from-id,length) pairs for each message. Same functionality as PetscGatherMessageLengths()
150   except it takes TWO ilenths and output TWO olengths.
151 
152   Collective on MPI_Comm
153 
154   Input Parameters:
155 + comm      - Communicator
156 . nsends    - number of messages that are to be sent.
157 . nrecvs    - number of messages being received
158 - ilengths1, ilengths2 - array of integers of length sizeof(comm)
159               a non zero ilengths[i] represent a message to i of length ilengths[i]
160 
161   Output Parameters:
162 + onodes    - list of node-ids from which messages are expected
163 - olengths1, olengths2 - corresponding message lengths
164 
165   Level: developer
166 
167   Concepts: mpi utility
168 
169   Notes:
170   With this info, the correct MPI_Irecv() can be posted with the correct
171   from-id, with a buffer with the right amount of memory required.
172 
173   The calling function deallocates the memory in onodes and olengths
174 
175   To determine nrecevs, one can use PetscGatherNumberOfMessages()
176 
177 .seealso: PetscGatherMessageLengths() and PetscGatherNumberOfMessages()
178 @*/
179 PetscErrorCode PETSC_DLLEXPORT PetscGatherMessageLengths2(MPI_Comm comm,PetscMPIInt nsends,PetscMPIInt nrecvs,const PetscMPIInt ilengths1[],const PetscMPIInt ilengths2[],PetscMPIInt **onodes,PetscMPIInt **olengths1,PetscMPIInt **olengths2)
180 {
181   PetscErrorCode ierr;
182   PetscMPIInt    size,tag,i,j,*buf_s = PETSC_NULL,*buf_r = PETSC_NULL,*buf_j = PETSC_NULL;
183   MPI_Request    *s_waits = PETSC_NULL,*r_waits = PETSC_NULL;
184   MPI_Status     *w_status = PETSC_NULL;
185 
186   PetscFunctionBegin;
187   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
188   ierr = PetscCommGetNewTag(comm,&tag);CHKERRQ(ierr);
189 
190   /* cannot use PetscMalloc5() because r_waits and s_waits must be contiquous for the call to MPI_Waitall() */
191   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);
192   s_waits = r_waits + nrecvs;
193 
194   /* Post the Irecv to get the message length-info */
195   ierr = PetscMalloc((nrecvs+1)*sizeof(PetscMPIInt),olengths1);CHKERRQ(ierr);
196   ierr = PetscMalloc((nrecvs+1)*sizeof(PetscMPIInt),olengths2);CHKERRQ(ierr);
197   for (i=0; i<nrecvs; i++) {
198     buf_j = buf_r + (2*i);
199     ierr = MPI_Irecv(buf_j,2,MPI_INT,MPI_ANY_SOURCE,tag,comm,r_waits+i);CHKERRQ(ierr);
200   }
201 
202   /* Post the Isends with the message length-info */
203   for (i=0,j=0; i<size; ++i) {
204     if (ilengths1[i]) {
205       buf_j = buf_s + (2*j);
206       buf_j[0] = *(ilengths1+i);
207       buf_j[1] = *(ilengths2+i);
208       ierr = MPI_Isend(buf_j,2,MPI_INT,i,tag,comm,s_waits+j);CHKERRQ(ierr);
209       j++;
210     }
211   }
212 
213   /* Post waits on sends and receivs */
214   if (nrecvs+nsends) {ierr = MPI_Waitall(nrecvs+nsends,r_waits,w_status);CHKERRQ(ierr);}
215 
216 
217   /* Pack up the received data */
218   ierr = PetscMalloc((nrecvs+1)*sizeof(PetscMPIInt),onodes);CHKERRQ(ierr);
219   for (i=0; i<nrecvs; ++i) {
220     (*onodes)[i] = w_status[i].MPI_SOURCE;
221     buf_j = buf_r + (2*i);
222     (*olengths1)[i] = buf_j[0];
223     (*olengths2)[i] = buf_j[1];
224   }
225 
226   ierr = PetscFree4(r_waits,buf_r,buf_s,w_status);CHKERRQ(ierr);
227   PetscFunctionReturn(0);
228 }
229 
230 /*
231 
232   Allocate a bufffer sufficient to hold messages of size specified in olengths.
233   And post Irecvs on these buffers using node info from onodes
234 
235  */
236 #undef __FUNCT__
237 #define __FUNCT__ "PetscPostIrecvInt"
238 PetscErrorCode PETSC_DLLEXPORT PetscPostIrecvInt(MPI_Comm comm,PetscMPIInt tag,PetscMPIInt nrecvs,const PetscMPIInt onodes[],const PetscMPIInt olengths[],PetscInt ***rbuf,MPI_Request **r_waits)
239 {
240   PetscErrorCode ierr;
241   PetscInt       **rbuf_t,i;
242   MPI_Request    *r_waits_t;
243   size_t         len = 0;
244 
245   PetscFunctionBegin;
246 
247   /* compute memory required for recv buffers */
248   for (i=0; i<nrecvs; i++) len += olengths[i];  /* each message length */
249   len *= sizeof(PetscInt);
250   len += (nrecvs+1)*sizeof(PetscInt*); /* Array of pointers for each message */
251 
252   /* allocate memory for recv buffers */
253   ierr    = PetscMalloc(len,&rbuf_t);CHKERRQ(ierr);
254   rbuf_t[0] = (PetscInt*)(rbuf_t + nrecvs);
255   for (i=1; i<nrecvs; ++i) rbuf_t[i] = rbuf_t[i-1] + olengths[i-1];
256 
257   /* Post the receives */
258   ierr = PetscMalloc((nrecvs+1)*sizeof(MPI_Request),&r_waits_t);CHKERRQ(ierr);
259   for (i=0; i<nrecvs; ++i) {
260     ierr = MPI_Irecv(rbuf_t[i],olengths[i],MPIU_INT,onodes[i],tag,comm,r_waits_t+i);CHKERRQ(ierr);
261   }
262 
263   *rbuf    = rbuf_t;
264   *r_waits = r_waits_t;
265   PetscFunctionReturn(0);
266 }
267 
268 #undef __FUNCT__
269 #define __FUNCT__ "PetscPostIrecvScalar"
270 PetscErrorCode PETSC_DLLEXPORT PetscPostIrecvScalar(MPI_Comm comm,PetscMPIInt tag,PetscMPIInt nrecvs,const PetscMPIInt onodes[],const PetscMPIInt olengths[],PetscScalar ***rbuf,MPI_Request **r_waits)
271 {
272   PetscErrorCode ierr;
273   PetscMPIInt    i;
274   PetscScalar    **rbuf_t;
275   MPI_Request    *r_waits_t;
276   size_t         len = 0;
277   PetscFunctionBegin;
278 
279   /* compute memory required for recv buffers */
280   for (i=0; i<nrecvs; i++) len += olengths[i];  /* each message length */
281   len *= sizeof(PetscScalar);
282   len += (nrecvs+1)*sizeof(PetscScalar*); /* Array of pointers for each message */
283 
284 
285   /* allocate memory for recv buffers */
286   ierr    = PetscMalloc(len,&rbuf_t);CHKERRQ(ierr);
287   rbuf_t[0] = (PetscScalar*)(rbuf_t + nrecvs);
288   for (i=1; i<nrecvs; ++i) rbuf_t[i] = rbuf_t[i-1] + olengths[i-1];
289 
290   /* Post the receives */
291   ierr = PetscMalloc((nrecvs+1)*sizeof(MPI_Request),&r_waits_t);CHKERRQ(ierr);
292   for (i=0; i<nrecvs; ++i) {
293     ierr = MPI_Irecv(rbuf_t[i],olengths[i],MPIU_SCALAR,onodes[i],tag,comm,r_waits_t+i);CHKERRQ(ierr);
294   }
295 
296   *rbuf    = rbuf_t;
297   *r_waits = r_waits_t;
298   PetscFunctionReturn(0);
299 }
300