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