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