xref: /petsc/src/sys/utils/mpimesg.c (revision 7c4f633dc6bb6149cca88d301ead35a99e103cbb)
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       len=0,**rbuf_t,i;
242   MPI_Request    *r_waits_t;
243 
244   PetscFunctionBegin;
245 
246   /* compute memory required for recv buffers */
247   for (i=0; i<nrecvs; i++) len += olengths[i];  /* each message length */
248   len *= sizeof(PetscInt);
249   len += (nrecvs+1)*sizeof(PetscInt*); /* Array of pointers for each message */
250 
251   /* allocate memory for recv buffers */
252   ierr    = PetscMalloc(len,&rbuf_t);CHKERRQ(ierr);
253   rbuf_t[0] = (PetscInt*)(rbuf_t + nrecvs);
254   for (i=1; i<nrecvs; ++i) rbuf_t[i] = rbuf_t[i-1] + olengths[i-1];
255 
256   /* Post the receives */
257   ierr = PetscMalloc((nrecvs+1)*sizeof(MPI_Request),&r_waits_t);CHKERRQ(ierr);
258   for (i=0; i<nrecvs; ++i) {
259     ierr = MPI_Irecv(rbuf_t[i],olengths[i],MPIU_INT,onodes[i],tag,comm,r_waits_t+i);CHKERRQ(ierr);
260   }
261 
262   *rbuf    = rbuf_t;
263   *r_waits = r_waits_t;
264   PetscFunctionReturn(0);
265 }
266 
267 #undef __FUNCT__
268 #define __FUNCT__ "PetscPostIrecvScalar"
269 PetscErrorCode PETSC_DLLEXPORT PetscPostIrecvScalar(MPI_Comm comm,PetscMPIInt tag,PetscMPIInt nrecvs,const PetscMPIInt onodes[],const PetscMPIInt olengths[],PetscScalar ***rbuf,MPI_Request **r_waits)
270 {
271   PetscErrorCode ierr;
272   PetscMPIInt    len=0,i;
273   PetscScalar    **rbuf_t;
274   MPI_Request    *r_waits_t;
275 
276   PetscFunctionBegin;
277 
278   /* compute memory required for recv buffers */
279   for (i=0; i<nrecvs; i++) len += olengths[i];  /* each message length */
280   len *= sizeof(PetscScalar);
281   len += (nrecvs+1)*sizeof(PetscScalar*); /* Array of pointers for each message */
282 
283 
284   /* allocate memory for recv buffers */
285   ierr    = PetscMalloc(len,&rbuf_t);CHKERRQ(ierr);
286   rbuf_t[0] = (PetscScalar*)(rbuf_t + nrecvs);
287   for (i=1; i<nrecvs; ++i) rbuf_t[i] = rbuf_t[i-1] + olengths[i-1];
288 
289   /* Post the receives */
290   ierr = PetscMalloc((nrecvs+1)*sizeof(MPI_Request),&r_waits_t);CHKERRQ(ierr);
291   for (i=0; i<nrecvs; ++i) {
292     ierr = MPI_Irecv(rbuf_t[i],olengths[i],MPIU_SCALAR,onodes[i],tag,comm,r_waits_t+i);CHKERRQ(ierr);
293   }
294 
295   *rbuf    = rbuf_t;
296   *r_waits = r_waits_t;
297   PetscFunctionReturn(0);
298 }
299