xref: /petsc/src/sys/utils/mpimesg.c (revision 48cdcd6b53bc91835c082adca72c212f70eee222)
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,&recv_buf,size,&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,&r_waits,nrecvs+nsends,&w_status);CHKERRQ(ierr);
113   s_waits = r_waits+nrecvs;
114 
115   /* Post the Irecv to get the message length-info */
116   ierr = PetscMalloc1(nrecvs,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 = PetscMalloc1(nrecvs,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,&r_waits,2*nrecvs,&buf_r,2*nsends,&buf_s,nrecvs+nsends,&w_status);CHKERRQ(ierr);
186   s_waits = r_waits + nrecvs;
187 
188   /* Post the Irecv to get the message length-info */
189   ierr = PetscMalloc1(nrecvs+1,olengths1);CHKERRQ(ierr);
190   ierr = PetscMalloc1(nrecvs+1,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   if (j != nsends) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_PLIB,"j %d not equal to expected number of sends %d\n",j,nsends);
207 
208   /* Post waits on sends and receivs */
209   if (nrecvs+nsends) {ierr = MPI_Waitall(nrecvs+nsends,r_waits,w_status);CHKERRQ(ierr);}
210 
211 
212   /* Pack up the received data */
213   ierr = PetscMalloc1(nrecvs+1,onodes);CHKERRQ(ierr);
214   for (i=0; i<nrecvs; ++i) {
215     (*onodes)[i]    = w_status[i].MPI_SOURCE;
216     buf_j           = buf_r + (2*i);
217     (*olengths1)[i] = buf_j[0];
218     (*olengths2)[i] = buf_j[1];
219   }
220 
221   ierr = PetscFree4(r_waits,buf_r,buf_s,w_status);CHKERRQ(ierr);
222   PetscFunctionReturn(0);
223 }
224 
225 /*
226 
227   Allocate a bufffer sufficient to hold messages of size specified in olengths.
228   And post Irecvs on these buffers using node info from onodes
229 
230  */
231 #undef __FUNCT__
232 #define __FUNCT__ "PetscPostIrecvInt"
233 PetscErrorCode  PetscPostIrecvInt(MPI_Comm comm,PetscMPIInt tag,PetscMPIInt nrecvs,const PetscMPIInt onodes[],const PetscMPIInt olengths[],PetscInt ***rbuf,MPI_Request **r_waits)
234 {
235   PetscErrorCode ierr;
236   PetscInt       **rbuf_t,i,len = 0;
237   MPI_Request    *r_waits_t;
238 
239   PetscFunctionBegin;
240   /* compute memory required for recv buffers */
241   for (i=0; i<nrecvs; i++) len += olengths[i];  /* each message length */
242 
243   /* allocate memory for recv buffers */
244   ierr = PetscMalloc1(nrecvs+1,&rbuf_t);CHKERRQ(ierr);
245   ierr = PetscMalloc1(len,&rbuf_t[0]);CHKERRQ(ierr);
246   for (i=1; i<nrecvs; ++i) rbuf_t[i] = rbuf_t[i-1] + olengths[i-1];
247 
248   /* Post the receives */
249   ierr = PetscMalloc1(nrecvs,&r_waits_t);CHKERRQ(ierr);
250   for (i=0; i<nrecvs; ++i) {
251     ierr = MPI_Irecv(rbuf_t[i],olengths[i],MPIU_INT,onodes[i],tag,comm,r_waits_t+i);CHKERRQ(ierr);
252   }
253 
254   *rbuf    = rbuf_t;
255   *r_waits = r_waits_t;
256   PetscFunctionReturn(0);
257 }
258 
259 #undef __FUNCT__
260 #define __FUNCT__ "PetscPostIrecvScalar"
261 PetscErrorCode  PetscPostIrecvScalar(MPI_Comm comm,PetscMPIInt tag,PetscMPIInt nrecvs,const PetscMPIInt onodes[],const PetscMPIInt olengths[],PetscScalar ***rbuf,MPI_Request **r_waits)
262 {
263   PetscErrorCode ierr;
264   PetscMPIInt    i;
265   PetscScalar    **rbuf_t;
266   MPI_Request    *r_waits_t;
267   PetscInt       len = 0;
268 
269   PetscFunctionBegin;
270   /* compute memory required for recv buffers */
271   for (i=0; i<nrecvs; i++) len += olengths[i];  /* each message length */
272 
273   /* allocate memory for recv buffers */
274   ierr = PetscMalloc1(nrecvs+1,&rbuf_t);CHKERRQ(ierr);
275   ierr = PetscMalloc1(len,&rbuf_t[0]);CHKERRQ(ierr);
276   for (i=1; i<nrecvs; ++i) rbuf_t[i] = rbuf_t[i-1] + olengths[i-1];
277 
278   /* Post the receives */
279   ierr = PetscMalloc1(nrecvs,&r_waits_t);CHKERRQ(ierr);
280   for (i=0; i<nrecvs; ++i) {
281     ierr = MPI_Irecv(rbuf_t[i],olengths[i],MPIU_SCALAR,onodes[i],tag,comm,r_waits_t+i);CHKERRQ(ierr);
282   }
283 
284   *rbuf    = rbuf_t;
285   *r_waits = r_waits_t;
286   PetscFunctionReturn(0);
287 }
288