xref: /petsc/src/vec/is/sf/impls/basic/sfbasic.c (revision 63f3c55c12ae2f190c391f6df6c540efe018a44a)
1 
2 #include <../src/vec/is/sf/impls/basic/sfbasic.h>
3 
4 /*===================================================================================*/
5 /*              Internal routines for PetscSFPack                              */
6 /*===================================================================================*/
7 
8 /* Return root and leaf MPI requests for communication in the given direction. If the requests have not been
9    initialized (since we use persistent requests), then initialize them.
10 */
11 static PetscErrorCode PetscSFPackGetReqs_Basic(PetscSF sf,PetscSFPack link,PetscSFDirection direction,MPI_Request **rootreqs,MPI_Request **leafreqs)
12 {
13   PetscErrorCode ierr;
14   PetscSF_Basic  *bas = (PetscSF_Basic*)sf->data;
15   PetscInt       i,j,nrootranks,ndrootranks,nleafranks,ndleafranks,disp;
16   const PetscInt *rootoffset,*leafoffset;
17   PetscMPIInt    n;
18   MPI_Comm       comm = PetscObjectComm((PetscObject)sf);
19   MPI_Datatype   unit = link->unit;
20   PetscMemType   rootmtype = link->rootmtype,leafmtype = link->leafmtype;
21 
22   PetscFunctionBegin;
23   if (rootreqs && !link->rootreqsinited[direction][rootmtype]) {
24     ierr = PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,NULL,&rootoffset,NULL);CHKERRQ(ierr);
25     if (direction == PETSCSF_LEAF2ROOT_REDUCE) {
26       for (i=ndrootranks,j=0; i<nrootranks; i++,j++) {
27         disp = (rootoffset[i] - rootoffset[ndrootranks])*link->unitbytes;
28         ierr = PetscMPIIntCast(rootoffset[i+1]-rootoffset[i],&n);CHKERRQ(ierr);
29         ierr = MPI_Recv_init(link->rootbuf[rootmtype]+disp,n,unit,bas->iranks[i],link->tag,comm,&link->rootreqs[direction][rootmtype][j]);CHKERRQ(ierr);
30       }
31     } else if (direction == PETSCSF_ROOT2LEAF_BCAST) {
32       for (i=ndrootranks,j=0; i<nrootranks; i++,j++) {
33         disp = (rootoffset[i] - rootoffset[ndrootranks])*link->unitbytes;
34         ierr = PetscMPIIntCast(rootoffset[i+1]-rootoffset[i],&n);CHKERRQ(ierr);
35         ierr = MPI_Send_init(link->rootbuf[rootmtype]+disp,n,unit,bas->iranks[i],link->tag,comm,&link->rootreqs[direction][rootmtype][j]);CHKERRQ(ierr);
36       }
37     } else SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Out-of-range PetscSFDirection = %d\n",(int)direction);
38     link->rootreqsinited[direction][rootmtype] = PETSC_TRUE;
39   }
40 
41   if (leafreqs && !link->leafreqsinited[direction][leafmtype]) {
42     ierr = PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,NULL,&leafoffset,NULL,NULL);CHKERRQ(ierr);
43     if (direction == PETSCSF_LEAF2ROOT_REDUCE) {
44       for (i=ndleafranks,j=0; i<nleafranks; i++,j++) {
45         disp = (leafoffset[i] - leafoffset[ndleafranks])*link->unitbytes;
46         ierr = PetscMPIIntCast(leafoffset[i+1]-leafoffset[i],&n);CHKERRQ(ierr);
47         ierr = MPI_Send_init(link->leafbuf[leafmtype]+disp,n,unit,sf->ranks[i],link->tag,comm,&link->leafreqs[direction][leafmtype][j]);CHKERRQ(ierr);
48       }
49     } else if (direction == PETSCSF_ROOT2LEAF_BCAST) {
50       for (i=ndleafranks,j=0; i<nleafranks; i++,j++) {
51         disp = (leafoffset[i] - leafoffset[ndleafranks])*link->unitbytes;
52         ierr = PetscMPIIntCast(leafoffset[i+1]-leafoffset[i],&n);CHKERRQ(ierr);
53         ierr = MPI_Recv_init(link->leafbuf[leafmtype]+disp,n,unit,sf->ranks[i],link->tag,comm,&link->leafreqs[direction][leafmtype][j]);CHKERRQ(ierr);
54       }
55     } else SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Out-of-range PetscSFDirection = %d\n",(int)direction);
56     link->leafreqsinited[direction][leafmtype] = PETSC_TRUE;
57   }
58 
59   if (rootreqs) *rootreqs = link->rootreqs[direction][rootmtype];
60   if (leafreqs) *leafreqs = link->leafreqs[direction][leafmtype];
61   PetscFunctionReturn(0);
62 }
63 
64 /* Common part shared by SFBasic and SFNeighbor based on the fact they all deal with sparse graphs. */
65 PETSC_INTERN PetscErrorCode PetscSFPackGet_Basic_Common(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,const void *leafdata,PetscInt nrootreqs,PetscInt nleafreqs,PetscSFPack *mylink)
66 {
67   PetscErrorCode    ierr;
68   PetscSF_Basic     *bas = (PetscSF_Basic*)sf->data;
69   PetscInt          i,j,nreqs,nrootranks,ndrootranks,nleafranks,ndleafranks;
70   const PetscInt    *rootoffset,*leafoffset;
71   PetscSFPack       *p,link;
72   PetscBool         match;
73 
74   PetscFunctionBegin;
75   ierr = PetscSFPackSetErrorOnUnsupportedOverlap(sf,unit,rootdata,leafdata);CHKERRQ(ierr);
76 
77   /* Look for types in cache */
78   for (p=&bas->avail; (link=*p); p=&link->next) {
79     ierr = MPIPetsc_Type_compare(unit,link->unit,&match);CHKERRQ(ierr);
80     if (match) {
81       *p = link->next; /* Remove from available list */
82       goto found;
83     }
84   }
85 
86   ierr = PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,NULL,&rootoffset,NULL);CHKERRQ(ierr);
87   ierr = PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,NULL,&leafoffset,NULL,NULL);CHKERRQ(ierr);
88   ierr = PetscNew(&link);CHKERRQ(ierr);
89   ierr = PetscSFPackSetUp_Host(sf,link,unit);CHKERRQ(ierr);
90   ierr = PetscCommGetNewTag(PetscObjectComm((PetscObject)sf),&link->tag);CHKERRQ(ierr); /* One tag per link */
91 
92   /* Allocate root, leaf, self buffers, and MPI requests */
93   link->rootbuflen = rootoffset[nrootranks]-rootoffset[ndrootranks];
94   link->leafbuflen = leafoffset[nleafranks]-leafoffset[ndleafranks];
95   link->selfbuflen = rootoffset[ndrootranks]*link->unitbytes;
96   link->nrootreqs  = nrootreqs;
97   link->nleafreqs  = nleafreqs;
98   nreqs            = (nrootreqs+nleafreqs)*4; /* Quadruple the requests since there are two communication directions and two memory types */
99   ierr             = PetscMalloc1(nreqs,&link->reqs);CHKERRQ(ierr);
100   for (i=0; i<nreqs; i++) link->reqs[i] = MPI_REQUEST_NULL; /* Initialized to NULL so that we know which need to be freed in Destroy */
101 
102   for (i=0; i<2; i++) { /* Two communication directions */
103     for (j=0; j<2; j++) { /* Two memory types */
104       link->rootreqs[i][j] = link->reqs + nrootreqs*(2*i+j);
105       link->leafreqs[i][j] = link->reqs + nrootreqs*4 + nleafreqs*(2*i+j);
106     }
107   }
108 
109 found:
110   link->rootmtype = rootmtype;
111   link->leafmtype = leafmtype;
112 #if defined(PETSC_HAVE_CUDA)
113   if (rootmtype == PETSC_MEMTYPE_DEVICE || leafmtype == PETSC_MEMTYPE_DEVICE) {ierr = PetscSFPackSetUp_Device(sf,link,unit);CHKERRQ(ierr);}
114 #endif
115   if (!link->rootbuf[rootmtype]) {ierr = PetscMallocWithMemType(rootmtype,link->rootbuflen*link->unitbytes,(void**)&link->rootbuf[rootmtype]);CHKERRQ(ierr);}
116   if (!link->leafbuf[leafmtype]) {ierr = PetscMallocWithMemType(leafmtype,link->leafbuflen*link->unitbytes,(void**)&link->leafbuf[leafmtype]);CHKERRQ(ierr);}
117   if (!link->selfbuf[rootmtype]) {ierr = PetscMallocWithMemType(rootmtype,link->selfbuflen*link->unitbytes,(void**)&link->selfbuf[rootmtype]);CHKERRQ(ierr);}
118   if (rootmtype != leafmtype && !link->selfbuf[leafmtype]) {ierr = PetscMallocWithMemType(leafmtype,link->selfbuflen*link->unitbytes,(void**)&link->selfbuf[leafmtype]);CHKERRQ(ierr);}
119   link->rkey = rootdata;
120   link->lkey = leafdata;
121   link->next = bas->inuse;
122   bas->inuse = link;
123 
124   *mylink    = link;
125   PetscFunctionReturn(0);
126 }
127 
128 static PetscErrorCode PetscSFPackGet_Basic(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,const void *leafdata,PetscSFDirection direction,PetscSFPack *mylink)
129 {
130   PetscErrorCode    ierr;
131   PetscInt          nrootranks,ndrootranks,nleafranks,ndleafranks;
132 
133   PetscFunctionBegin;
134   ierr = PetscSFGetRootInfo_Basic(sf,&nrootranks,&ndrootranks,NULL,NULL,NULL);CHKERRQ(ierr);
135   ierr = PetscSFGetLeafInfo_Basic(sf,&nleafranks,&ndleafranks,NULL,NULL,NULL,NULL);CHKERRQ(ierr);
136   ierr = PetscSFPackGet_Basic_Common(sf,unit,rootmtype,rootdata,leafmtype,leafdata,nrootranks-ndrootranks,nleafranks-ndleafranks,mylink);CHKERRQ(ierr);
137   PetscFunctionReturn(0);
138 }
139 
140 /*===================================================================================*/
141 /*              SF public interface implementations                                  */
142 /*===================================================================================*/
143 PETSC_INTERN PetscErrorCode PetscSFSetUp_Basic(PetscSF sf)
144 {
145   PetscErrorCode ierr;
146   PetscSF_Basic  *bas = (PetscSF_Basic*)sf->data;
147   PetscInt       *rlengths,*ilengths,i;
148   PetscMPIInt    rank,niranks,*iranks,tag;
149   MPI_Comm       comm;
150   MPI_Group      group;
151   MPI_Request    *rootreqs,*leafreqs;
152 
153   PetscFunctionBegin;
154   ierr = MPI_Comm_group(PETSC_COMM_SELF,&group);CHKERRQ(ierr);
155   ierr = PetscSFSetUpRanks(sf,group);CHKERRQ(ierr);
156   ierr = MPI_Group_free(&group);CHKERRQ(ierr);
157   ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr);
158   ierr = PetscObjectGetNewTag((PetscObject)sf,&tag);CHKERRQ(ierr);
159   ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
160   /*
161    * Inform roots about how many leaves and from which ranks
162    */
163   ierr = PetscMalloc1(sf->nranks,&rlengths);CHKERRQ(ierr);
164   /* Determine number, sending ranks, and length of incoming */
165   for (i=0; i<sf->nranks; i++) {
166     rlengths[i] = sf->roffset[i+1] - sf->roffset[i]; /* Number of roots referenced by my leaves; for rank sf->ranks[i] */
167   }
168   ierr = PetscCommBuildTwoSided(comm,1,MPIU_INT,sf->nranks-sf->ndranks,sf->ranks+sf->ndranks,rlengths+sf->ndranks,&niranks,&iranks,(void**)&ilengths);CHKERRQ(ierr);
169 
170   /* Sort iranks. See use of VecScatterGetRemoteOrdered_Private() in MatGetBrowsOfAoCols_MPIAIJ() on why.
171      We could sort ranks there at the price of allocating extra working arrays. Presumably, niranks is
172      small and the sorting is cheap.
173    */
174   ierr = PetscSortMPIIntWithIntArray(niranks,iranks,ilengths);CHKERRQ(ierr);
175 
176   /* Partition into distinguished and non-distinguished incoming ranks */
177   bas->ndiranks = sf->ndranks;
178   bas->niranks = bas->ndiranks + niranks;
179   ierr = PetscMalloc2(bas->niranks,&bas->iranks,bas->niranks+1,&bas->ioffset);CHKERRQ(ierr);
180   bas->ioffset[0] = 0;
181   for (i=0; i<bas->ndiranks; i++) {
182     bas->iranks[i] = sf->ranks[i];
183     bas->ioffset[i+1] = bas->ioffset[i] + rlengths[i];
184   }
185   if (bas->ndiranks > 1 || (bas->ndiranks == 1 && bas->iranks[0] != rank)) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Broken setup for shared ranks");
186   for ( ; i<bas->niranks; i++) {
187     bas->iranks[i] = iranks[i-bas->ndiranks];
188     bas->ioffset[i+1] = bas->ioffset[i] + ilengths[i-bas->ndiranks];
189   }
190   bas->itotal = bas->ioffset[i];
191   ierr = PetscFree(rlengths);CHKERRQ(ierr);
192   ierr = PetscFree(iranks);CHKERRQ(ierr);
193   ierr = PetscFree(ilengths);CHKERRQ(ierr);
194 
195   /* Send leaf identities to roots */
196   ierr = PetscMalloc1(bas->itotal,&bas->irootloc);CHKERRQ(ierr);
197   ierr = PetscMalloc2(bas->niranks-bas->ndiranks,&rootreqs,sf->nranks-sf->ndranks,&leafreqs);CHKERRQ(ierr);
198   for (i=bas->ndiranks; i<bas->niranks; i++) {
199     ierr = MPI_Irecv(bas->irootloc+bas->ioffset[i],bas->ioffset[i+1]-bas->ioffset[i],MPIU_INT,bas->iranks[i],tag,comm,&rootreqs[i-bas->ndiranks]);CHKERRQ(ierr);
200   }
201   for (i=0; i<sf->nranks; i++) {
202     PetscMPIInt npoints;
203     ierr = PetscMPIIntCast(sf->roffset[i+1] - sf->roffset[i],&npoints);CHKERRQ(ierr);
204     if (i < sf->ndranks) {
205       if (sf->ranks[i] != rank) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Cannot interpret distinguished leaf rank");
206       if (bas->iranks[0] != rank) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Cannot interpret distinguished root rank");
207       if (npoints != bas->ioffset[1]-bas->ioffset[0]) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Distinguished rank exchange has mismatched lengths");
208       ierr = PetscArraycpy(bas->irootloc+bas->ioffset[0],sf->rremote+sf->roffset[i],npoints);CHKERRQ(ierr);
209       continue;
210     }
211     ierr = MPI_Isend(sf->rremote+sf->roffset[i],npoints,MPIU_INT,sf->ranks[i],tag,comm,&leafreqs[i-sf->ndranks]);CHKERRQ(ierr);
212   }
213   ierr = MPI_Waitall(bas->niranks-bas->ndiranks,rootreqs,MPI_STATUSES_IGNORE);CHKERRQ(ierr);
214   ierr = MPI_Waitall(sf->nranks-sf->ndranks,leafreqs,MPI_STATUSES_IGNORE);CHKERRQ(ierr);
215   ierr = PetscFree2(rootreqs,leafreqs);CHKERRQ(ierr);
216 
217   sf->selfleafdups    = PETSC_TRUE; /* The conservative assumption is there are data race */
218   sf->remoteleafdups  = PETSC_TRUE;
219   bas->selfrootdups   = PETSC_TRUE;
220   bas->remoterootdups = PETSC_TRUE;
221 
222   /* Setup packing optimization for roots and leaves */
223   ierr = PetscSFPackSetupOptimizations_Basic(sf);CHKERRQ(ierr);
224   PetscFunctionReturn(0);
225 }
226 
227 static PetscErrorCode PetscSFSetFromOptions_Basic(PetscOptionItems *PetscOptionsObject,PetscSF sf)
228 {
229   PetscErrorCode ierr;
230 
231   PetscFunctionBegin;
232   ierr = PetscOptionsHead(PetscOptionsObject,"PetscSF Basic options");CHKERRQ(ierr);
233   ierr = PetscOptionsTail();CHKERRQ(ierr);
234   PetscFunctionReturn(0);
235 }
236 
237 PETSC_INTERN PetscErrorCode PetscSFReset_Basic(PetscSF sf)
238 {
239   PetscSF_Basic     *bas = (PetscSF_Basic*)sf->data;
240   PetscErrorCode    ierr;
241 
242   PetscFunctionBegin;
243   if (bas->inuse) SETERRQ(PetscObjectComm((PetscObject)sf),PETSC_ERR_ARG_WRONGSTATE,"Outstanding operation has not been completed");
244   ierr = PetscFree2(bas->iranks,bas->ioffset);CHKERRQ(ierr);
245   ierr = PetscFree(bas->irootloc);CHKERRQ(ierr);
246 #if defined(PETSC_HAVE_CUDA)
247   if (bas->irootloc_d) {cudaError_t err = cudaFree(bas->irootloc_d);CHKERRCUDA(err);bas->irootloc_d=NULL;}
248 #endif
249   ierr = PetscSFPackDestroyAvailable(&bas->avail);CHKERRQ(ierr);
250   ierr = PetscSFPackDestroyOptimizations_Basic(sf);CHKERRQ(ierr);
251   PetscFunctionReturn(0);
252 }
253 
254 PETSC_INTERN PetscErrorCode PetscSFDestroy_Basic(PetscSF sf)
255 {
256   PetscErrorCode ierr;
257 
258   PetscFunctionBegin;
259   ierr = PetscSFReset(sf);CHKERRQ(ierr);
260   ierr = PetscFree(sf->data);CHKERRQ(ierr);
261   PetscFunctionReturn(0);
262 }
263 
264 PETSC_INTERN PetscErrorCode PetscSFView_Basic(PetscSF sf,PetscViewer viewer)
265 {
266   PetscErrorCode ierr;
267   PetscBool      iascii;
268 
269   PetscFunctionBegin;
270   ierr = PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);CHKERRQ(ierr);
271   if (iascii) {ierr = PetscViewerASCIIPrintf(viewer,"  sort=%s\n",sf->rankorder ? "rank-order" : "unordered");CHKERRQ(ierr);}
272   PetscFunctionReturn(0);
273 }
274 
275 static PetscErrorCode PetscSFBcastAndOpBegin_Basic(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op)
276 {
277   PetscErrorCode    ierr;
278   PetscSFPack       link;
279   const PetscInt    *rootloc = NULL;
280   MPI_Request       *rootreqs,*leafreqs;
281 
282   PetscFunctionBegin;
283   ierr = PetscSFPackGet_Basic(sf,unit,rootmtype,rootdata,leafmtype,leafdata,PETSCSF_ROOT2LEAF_BCAST,&link);CHKERRQ(ierr);
284   ierr = PetscSFGetRootIndicesWithMemType_Basic(sf,rootmtype,&rootloc);CHKERRQ(ierr);
285 
286   ierr = PetscSFPackGetReqs_Basic(sf,link,PETSCSF_ROOT2LEAF_BCAST,&rootreqs,&leafreqs);CHKERRQ(ierr);
287   /* Post Irecv. Note distinguished ranks receive data via shared memory (i.e., not via MPI) */
288   ierr = MPI_Startall_irecv(link->leafbuflen,unit,link->nleafreqs,leafreqs);CHKERRQ(ierr);
289 
290   /* Do Isend */
291   ierr = PetscSFPackRootData(sf,link,rootloc,rootdata,PETSC_TRUE);CHKERRQ(ierr);
292   ierr = MPI_Startall_isend(link->rootbuflen,unit,link->nrootreqs,rootreqs);CHKERRQ(ierr);
293 
294   /* Do self to self communication via memcpy only when rootdata and leafdata are in different memory */
295   if (rootmtype != leafmtype) {ierr = PetscMemcpyWithMemType(leafmtype,rootmtype,link->selfbuf[leafmtype],link->selfbuf[rootmtype],link->selfbuflen*link->unitbytes);CHKERRQ(ierr);}
296   PetscFunctionReturn(0);
297 }
298 
299 PETSC_INTERN PetscErrorCode PetscSFBcastAndOpEnd_Basic(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op)
300 {
301   PetscErrorCode    ierr;
302   PetscSFPack       link;
303   const PetscInt    *leafloc = NULL;
304 
305   PetscFunctionBegin;
306   ierr = PetscSFPackGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
307   ierr = PetscSFPackWaitall_Basic(link,PETSCSF_ROOT2LEAF_BCAST);CHKERRQ(ierr);
308   ierr = PetscSFGetLeafIndicesWithMemType_Basic(sf,leafmtype,&leafloc);CHKERRQ(ierr);
309   ierr = PetscSFUnpackAndOpLeafData(sf,link,leafloc,leafdata,op,PETSC_TRUE);CHKERRQ(ierr);
310   ierr = PetscSFPackReclaim(sf,&link);CHKERRQ(ierr);
311   PetscFunctionReturn(0);
312 }
313 
314 /* leaf -> root with reduction */
315 static PetscErrorCode PetscSFReduceBegin_Basic(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op)
316 {
317   PetscErrorCode    ierr;
318   PetscSFPack       link;
319   const PetscInt    *leafloc = NULL;
320   MPI_Request       *rootreqs = NULL,*leafreqs = NULL; /* dummy values for compiler warnings about uninitialized value */
321 
322   PetscFunctionBegin;
323   ierr = PetscSFGetLeafIndicesWithMemType_Basic(sf,leafmtype,&leafloc);
324 
325   ierr = PetscSFPackGet_Basic(sf,unit,rootmtype,rootdata,leafmtype,leafdata,PETSCSF_LEAF2ROOT_REDUCE,&link);CHKERRQ(ierr);
326   ierr = PetscSFPackGetReqs_Basic(sf,link,PETSCSF_LEAF2ROOT_REDUCE,&rootreqs,&leafreqs);CHKERRQ(ierr);
327   /* Eagerly post root receives for non-distinguished ranks */
328   ierr = MPI_Startall_irecv(link->rootbuflen,unit,link->nrootreqs,rootreqs);CHKERRQ(ierr);
329 
330   /* Pack and send leaf data */
331   ierr = PetscSFPackLeafData(sf,link,leafloc,leafdata,PETSC_TRUE);CHKERRQ(ierr);
332   ierr = MPI_Startall_isend(link->leafbuflen,unit,link->nleafreqs,leafreqs);CHKERRQ(ierr);
333 
334   if (rootmtype != leafmtype) {ierr = PetscMemcpyWithMemType(rootmtype,leafmtype,link->selfbuf[rootmtype],link->selfbuf[leafmtype],link->selfbuflen*link->unitbytes);CHKERRQ(ierr);}
335   PetscFunctionReturn(0);
336 }
337 
338 PETSC_INTERN PetscErrorCode PetscSFReduceEnd_Basic(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op)
339 {
340   PetscErrorCode    ierr;
341   PetscSFPack       link;
342   const PetscInt    *rootloc = NULL;
343 
344   PetscFunctionBegin;
345   ierr = PetscSFGetRootIndicesWithMemType_Basic(sf,rootmtype,&rootloc);CHKERRQ(ierr);
346   ierr = PetscSFPackGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
347   ierr = PetscSFPackWaitall_Basic(link,PETSCSF_LEAF2ROOT_REDUCE);CHKERRQ(ierr);
348   ierr = PetscSFUnpackAndOpRootData(sf,link,rootloc,rootdata,op,PETSC_TRUE);CHKERRQ(ierr);
349   ierr = PetscSFPackReclaim(sf,&link);CHKERRQ(ierr);
350   PetscFunctionReturn(0);
351 }
352 
353 PETSC_INTERN PetscErrorCode PetscSFFetchAndOpBegin_Basic(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,void *rootdata,PetscMemType leafmtype,const void *leafdata,void *leafupdate,MPI_Op op)
354 {
355   PetscErrorCode ierr;
356 
357   PetscFunctionBegin;
358   ierr = PetscSFReduceBegin(sf,unit,leafdata,rootdata,op);CHKERRQ(ierr);
359   PetscFunctionReturn(0);
360 }
361 
362 static PetscErrorCode PetscSFFetchAndOpEnd_Basic(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,void *rootdata,PetscMemType leafmtype,const void *leafdata,void *leafupdate,MPI_Op op)
363 {
364   PetscErrorCode    ierr;
365   PetscSFPack       link;
366   const PetscInt    *rootloc = NULL,*leafloc = NULL;
367   MPI_Request       *rootreqs,*leafreqs;
368 
369   PetscFunctionBegin;
370   ierr = PetscSFGetRootIndicesWithMemType_Basic(sf,rootmtype,&rootloc);CHKERRQ(ierr);
371   ierr = PetscSFGetLeafIndicesWithMemType_Basic(sf,leafmtype,&leafloc);CHKERRQ(ierr);
372   ierr = PetscSFPackGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
373   /* This implementation could be changed to unpack as receives arrive, at the cost of non-determinism */
374   ierr = PetscSFPackWaitall_Basic(link,PETSCSF_LEAF2ROOT_REDUCE);CHKERRQ(ierr);
375   ierr = PetscSFPackGetReqs_Basic(sf,link,PETSCSF_ROOT2LEAF_BCAST,&rootreqs,&leafreqs);CHKERRQ(ierr);
376 
377   /* Post leaf receives */
378   ierr = MPI_Startall_irecv(link->leafbuflen,unit,link->nleafreqs,leafreqs);CHKERRQ(ierr);
379 
380   /* Process local fetch-and-op, post root sends */
381   ierr = PetscSFFetchAndOpRootData(sf,link,rootloc,rootdata,op,PETSC_TRUE);CHKERRQ(ierr);
382   ierr = MPI_Startall_isend(link->rootbuflen,unit,link->nrootreqs,rootreqs);CHKERRQ(ierr);
383   if (rootmtype != leafmtype) {ierr = PetscMemcpyWithMemType(leafmtype,rootmtype,link->selfbuf[leafmtype],link->selfbuf[rootmtype],link->selfbuflen*link->unitbytes);CHKERRQ(ierr);}
384 
385   /* Unpack and insert fetched data into leaves */
386   ierr = PetscSFPackWaitall_Basic(link,PETSCSF_ROOT2LEAF_BCAST);CHKERRQ(ierr);
387   ierr = PetscSFUnpackAndOpLeafData(sf,link,leafloc,leafupdate,MPIU_REPLACE,PETSC_TRUE);CHKERRQ(ierr);
388   ierr = PetscSFPackReclaim(sf,&link);CHKERRQ(ierr);
389   PetscFunctionReturn(0);
390 }
391 
392 PETSC_INTERN PetscErrorCode PetscSFGetLeafRanks_Basic(PetscSF sf,PetscInt *niranks,const PetscMPIInt **iranks,const PetscInt **ioffset,const PetscInt **irootloc)
393 {
394   PetscSF_Basic *bas = (PetscSF_Basic*)sf->data;
395 
396   PetscFunctionBegin;
397   if (niranks)  *niranks  = bas->niranks;
398   if (iranks)   *iranks   = bas->iranks;
399   if (ioffset)  *ioffset  = bas->ioffset;
400   if (irootloc) *irootloc = bas->irootloc;
401   PetscFunctionReturn(0);
402 }
403 
404 /* An optimized PetscSFCreateEmbeddedSF. We aggresively make use of the established communication on sf.
405    We need one bcast on sf, and no communication anymore to build the embedded sf. Note that selected[]
406    was sorted before calling the routine.
407  */
408 PETSC_INTERN PetscErrorCode PetscSFCreateEmbeddedSF_Basic(PetscSF sf,PetscInt nselected,const PetscInt *selected,PetscSF *newsf)
409 {
410   PetscSF           esf;
411   PetscInt          esf_nranks,esf_ndranks,*esf_roffset,*esf_rmine,*esf_rremote,count;
412   PetscInt          i,j,k,p,q,nroots,*rootdata,*leafdata,connected_leaves,*new_ilocal,nranks,ndranks,niranks,ndiranks,minleaf,maxleaf,maxlocal;
413   PetscMPIInt       *esf_ranks;
414   const PetscMPIInt *ranks,*iranks;
415   const PetscInt    *roffset,*rmine,*rremote,*ioffset,*irootloc,*buffer;
416   PetscBool         connected;
417   PetscSFPack       link;
418   PetscSFNode       *new_iremote;
419   PetscSF_Basic     *bas;
420   PetscErrorCode    ierr;
421 
422   PetscFunctionBegin;
423   ierr = PetscSFCreate(PetscObjectComm((PetscObject)sf),&esf);CHKERRQ(ierr);
424   ierr = PetscSFSetType(esf,PETSCSFBASIC);CHKERRQ(ierr); /* This optimized routine can only create a basic sf */
425 
426   /* Find out which leaves are still connected to roots in the embedded sf */
427   ierr = PetscSFGetGraph(sf,&nroots,NULL,NULL,NULL);CHKERRQ(ierr);
428   ierr = PetscSFGetLeafRange(sf,&minleaf,&maxleaf);CHKERRQ(ierr);
429   /* We abused the term leafdata here, whose size is usually the number of leaf data items. Here its size is # of leaves (always >= # of leaf data items) */
430   maxlocal = (minleaf > maxleaf)? 0 : maxleaf-minleaf+1; /* maxleaf=-1 and minleaf=0 when nleaves=0 */
431   ierr = PetscCalloc2(nroots,&rootdata,maxlocal,&leafdata);CHKERRQ(ierr);
432   /* Tag selected roots */
433   for (i=0; i<nselected; ++i) rootdata[selected[i]] = 1;
434 
435   /* Bcast from roots to leaves to tag connected leaves. We reuse the established bcast communication in
436      sf but do not do unpacking (from leaf buffer to leafdata). The raw data in leaf buffer is what we are
437      interested in since it tells which leaves are connected to which ranks.
438    */
439   ierr = PetscSFBcastAndOpBegin_Basic(sf,MPIU_INT,PETSC_MEMTYPE_HOST,rootdata,PETSC_MEMTYPE_HOST,leafdata-minleaf,MPIU_REPLACE);CHKERRQ(ierr); /* Need to give leafdata but we won't use it */
440   ierr = PetscSFPackGetInUse(sf,MPIU_INT,rootdata,leafdata-minleaf,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
441   ierr = PetscSFPackWaitall_Basic(link,PETSCSF_ROOT2LEAF_BCAST);CHKERRQ(ierr);
442   ierr = PetscSFGetLeafInfo_Basic(sf,&nranks,&ndranks,&ranks,&roffset,&rmine,&rremote);CHKERRQ(ierr); /* Get send info */
443   esf_nranks = esf_ndranks = connected_leaves = 0;
444   for (i=0; i<nranks; i++) {
445     connected = PETSC_FALSE; /* Is the current process still connected to this remote root rank? */
446     buffer    = i < ndranks? (PetscInt*)link->selfbuf[PETSC_MEMTYPE_HOST] : (PetscInt*)link->leafbuf[PETSC_MEMTYPE_HOST] + (roffset[i] - roffset[ndranks]);
447     count     = roffset[i+1] - roffset[i];
448     for (j=0; j<count; j++) {if (buffer[j]) {connected_leaves++; connected = PETSC_TRUE;}}
449     if (connected) {esf_nranks++; if (i < ndranks) esf_ndranks++;}
450   }
451 
452   /* Set graph of esf and also set up its outgoing communication (i.e., send info), which is usually done by PetscSFSetUpRanks */
453   ierr = PetscMalloc1(connected_leaves,&new_ilocal);CHKERRQ(ierr);
454   ierr = PetscMalloc1(connected_leaves,&new_iremote);CHKERRQ(ierr);
455   ierr = PetscMalloc4(esf_nranks,&esf_ranks,esf_nranks+1,&esf_roffset,connected_leaves,&esf_rmine,connected_leaves,&esf_rremote);CHKERRQ(ierr);
456   p    = 0; /* Counter for connected root ranks */
457   q    = 0; /* Counter for connected leaves */
458   esf_roffset[0] = 0;
459   for (i=0; i<nranks; i++) { /* Scan leaf data again to fill esf arrays */
460     buffer    = i < ndranks? (PetscInt*)link->selfbuf[PETSC_MEMTYPE_HOST] : (PetscInt*)link->leafbuf[PETSC_MEMTYPE_HOST] + (roffset[i] - roffset[ndranks]);
461     connected = PETSC_FALSE;
462     for (j=roffset[i],k=0; j<roffset[i+1]; j++,k++) {
463         if (buffer[k]) {
464         esf_rmine[q]         = new_ilocal[q] = rmine[j];
465         esf_rremote[q]       = rremote[j];
466         new_iremote[q].index = rremote[j];
467         new_iremote[q].rank  = ranks[i];
468         connected            = PETSC_TRUE;
469         q++;
470       }
471     }
472     if (connected) {
473       esf_ranks[p]     = ranks[i];
474       esf_roffset[p+1] = q;
475       p++;
476     }
477   }
478 
479   ierr = PetscSFPackReclaim(sf,&link);CHKERRQ(ierr);
480 
481   /* SetGraph internally resets the SF, so we only set its fields after the call */
482   ierr         = PetscSFSetGraph(esf,nroots,connected_leaves,new_ilocal,PETSC_OWN_POINTER,new_iremote,PETSC_OWN_POINTER);CHKERRQ(ierr);
483   esf->nranks  = esf_nranks;
484   esf->ndranks = esf_ndranks;
485   esf->ranks   = esf_ranks;
486   esf->roffset = esf_roffset;
487   esf->rmine   = esf_rmine;
488   esf->rremote = esf_rremote;
489 
490   /* Set up the incoming communication (i.e., recv info) stored in esf->data, which is usually done by PetscSFSetUp_Basic */
491   bas  = (PetscSF_Basic*)esf->data;
492   ierr = PetscSFGetRootInfo_Basic(sf,&niranks,&ndiranks,&iranks,&ioffset,&irootloc);CHKERRQ(ierr); /* Get recv info */
493   /* Embedded sf always has simpler communication than the original one. We might allocate longer arrays than needed here. But we
494      expect these arrays are usually short, so we do not care. The benefit is we can fill these arrays by just parsing irootloc once.
495    */
496   ierr = PetscMalloc2(niranks,&bas->iranks,niranks+1,&bas->ioffset);CHKERRQ(ierr);
497   ierr = PetscMalloc1(ioffset[niranks],&bas->irootloc);CHKERRQ(ierr);
498   bas->niranks = bas->ndiranks = bas->ioffset[0] = 0;
499   p = 0; /* Counter for connected leaf ranks */
500   q = 0; /* Counter for connected roots */
501   for (i=0; i<niranks; i++) {
502     connected = PETSC_FALSE; /* Is the current process still connected to this remote leaf rank? */
503     for (j=ioffset[i]; j<ioffset[i+1]; j++) {
504       PetscInt loc;
505       ierr = PetscFindInt(irootloc[j],nselected,selected,&loc);CHKERRQ(ierr);
506       if (loc >= 0) { /* Found in selected this root is connected */
507         bas->irootloc[q++] = irootloc[j];
508         connected = PETSC_TRUE;
509       }
510     }
511     if (connected) {
512       bas->niranks++;
513       if (i<ndiranks) bas->ndiranks++; /* Note that order of ranks (including distinguished ranks) is kept */
514       bas->iranks[p]    = iranks[i];
515       bas->ioffset[p+1] = q;
516       p++;
517     }
518   }
519   bas->itotal = q;
520 
521   /* Setup packing optimizations */
522   ierr = PetscSFPackSetupOptimizations_Basic(esf);CHKERRQ(ierr);
523   esf->setupcalled = PETSC_TRUE; /* We have done setup ourselves! */
524 
525   ierr = PetscFree2(rootdata,leafdata);CHKERRQ(ierr);
526   *newsf = esf;
527   PetscFunctionReturn(0);
528 }
529 
530 PETSC_INTERN PetscErrorCode PetscSFCreateEmbeddedLeafSF_Basic(PetscSF sf,PetscInt nselected,const PetscInt *selected,PetscSF *newsf)
531 {
532   PetscSF           esf;
533   PetscInt          i,j,k,p,q,nroots,*rootdata,*leafdata,*new_ilocal,niranks,ndiranks,minleaf,maxleaf,maxlocal;
534   const PetscInt    *ilocal,*ioffset,*irootloc,*buffer;
535   const PetscMPIInt *iranks;
536   PetscSFPack       link;
537   PetscSFNode       *new_iremote;
538   const PetscSFNode *iremote;
539   PetscSF_Basic     *bas;
540   MPI_Group         group;
541   PetscErrorCode    ierr;
542 
543   PetscFunctionBegin;
544   ierr = PetscSFCreate(PetscObjectComm((PetscObject)sf),&esf);CHKERRQ(ierr);
545   ierr = PetscSFSetType(esf,PETSCSFBASIC);CHKERRQ(ierr); /* This optimized routine can only create a basic sf */
546 
547   /* Set the graph of esf, which is easy for CreateEmbeddedLeafSF */
548   ierr = PetscSFGetGraph(sf,&nroots,NULL,&ilocal,&iremote);CHKERRQ(ierr);
549   ierr = PetscSFGetLeafRange(sf,&minleaf,&maxleaf);CHKERRQ(ierr);
550   ierr = PetscMalloc1(nselected,&new_ilocal);CHKERRQ(ierr);
551   ierr = PetscMalloc1(nselected,&new_iremote);CHKERRQ(ierr);
552   for (i=0; i<nselected; i++) {
553     const PetscInt l     = selected[i];
554     new_ilocal[i]        = ilocal ? ilocal[l] : l;
555     new_iremote[i].rank  = iremote[l].rank;
556     new_iremote[i].index = iremote[l].index;
557   }
558 
559   /* Tag selected leaves before PetscSFSetGraph since new_ilocal might turn into NULL since we use PETSC_OWN_POINTER below */
560   maxlocal = (minleaf > maxleaf)? 0 : maxleaf-minleaf+1; /* maxleaf=-1 and minleaf=0 when nleaves=0 */
561   ierr = PetscCalloc2(nroots,&rootdata,maxlocal,&leafdata);CHKERRQ(ierr);
562   for (i=0; i<nselected; i++) leafdata[new_ilocal[i]-minleaf] = 1; /* -minleaf to adjust indices according to minleaf */
563 
564   ierr = PetscSFSetGraph(esf,nroots,nselected,new_ilocal,PETSC_OWN_POINTER,new_iremote,PETSC_OWN_POINTER);CHKERRQ(ierr);
565 
566   /* Set up the outgoing communication (i.e., send info). We can not reuse rmine etc in sf since there is no way to
567      map rmine[i] (ilocal of leaves) back to selected[j]  (leaf indices).
568    */
569   ierr = MPI_Comm_group(PETSC_COMM_SELF,&group);CHKERRQ(ierr);
570   ierr = PetscSFSetUpRanks(esf,group);CHKERRQ(ierr);
571   ierr = MPI_Group_free(&group);CHKERRQ(ierr);
572 
573   /* Set up the incoming communication (i.e., recv info) */
574   ierr = PetscSFGetRootInfo_Basic(sf,&niranks,&ndiranks,&iranks,&ioffset,&irootloc);CHKERRQ(ierr);
575   bas  = (PetscSF_Basic*)esf->data;
576   ierr = PetscMalloc2(niranks,&bas->iranks,niranks+1,&bas->ioffset);CHKERRQ(ierr);
577   ierr = PetscMalloc1(ioffset[niranks],&bas->irootloc);CHKERRQ(ierr);
578 
579   /* Pass info about selected leaves to root buffer */
580   ierr = PetscSFReduceBegin_Basic(sf,MPIU_INT,PETSC_MEMTYPE_HOST,leafdata-minleaf,PETSC_MEMTYPE_HOST,rootdata,MPIU_REPLACE);CHKERRQ(ierr); /* -minleaf to re-adjust start address of leafdata */
581   ierr = PetscSFPackGetInUse(sf,MPIU_INT,rootdata,leafdata-minleaf,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
582   ierr = PetscSFPackWaitall_Basic(link,PETSCSF_LEAF2ROOT_REDUCE);CHKERRQ(ierr);
583 
584   bas->niranks = bas->ndiranks = bas->ioffset[0] = 0;
585   p = 0; /* Counter for connected leaf ranks */
586   q = 0; /* Counter for connected roots */
587   for (i=0; i<niranks; i++) {
588     PetscBool connected = PETSC_FALSE; /* Is the current process still connected to this remote leaf rank? */
589     buffer = i < ndiranks? (PetscInt*)link->selfbuf[PETSC_MEMTYPE_HOST] : (PetscInt*)link->rootbuf[PETSC_MEMTYPE_HOST] + (ioffset[i] - ioffset[ndiranks]);
590     for (j=ioffset[i],k=0; j<ioffset[i+1]; j++,k++) {
591       if (buffer[k]) {bas->irootloc[q++] = irootloc[j]; connected = PETSC_TRUE;}
592     }
593     if (connected) {
594       bas->niranks++;
595       if (i<ndiranks) bas->ndiranks++;
596       bas->iranks[p]    = iranks[i];
597       bas->ioffset[p+1] = q;
598       p++;
599     }
600   }
601   bas->itotal = q;
602   ierr = PetscSFPackReclaim(sf,&link);CHKERRQ(ierr);
603 
604   /* Setup packing optimizations */
605   ierr = PetscSFPackSetupOptimizations_Basic(esf);CHKERRQ(ierr);
606   esf->setupcalled = PETSC_TRUE; /* We have done setup ourselves! */
607 
608   ierr = PetscFree2(rootdata,leafdata);CHKERRQ(ierr);
609   *newsf = esf;
610   PetscFunctionReturn(0);
611 }
612 
613 PETSC_EXTERN PetscErrorCode PetscSFCreate_Basic(PetscSF sf)
614 {
615   PetscSF_Basic  *dat;
616   PetscErrorCode ierr;
617 
618   PetscFunctionBegin;
619   sf->ops->SetUp                = PetscSFSetUp_Basic;
620   sf->ops->SetFromOptions       = PetscSFSetFromOptions_Basic;
621   sf->ops->Reset                = PetscSFReset_Basic;
622   sf->ops->Destroy              = PetscSFDestroy_Basic;
623   sf->ops->View                 = PetscSFView_Basic;
624   sf->ops->BcastAndOpBegin      = PetscSFBcastAndOpBegin_Basic;
625   sf->ops->BcastAndOpEnd        = PetscSFBcastAndOpEnd_Basic;
626   sf->ops->ReduceBegin          = PetscSFReduceBegin_Basic;
627   sf->ops->ReduceEnd            = PetscSFReduceEnd_Basic;
628   sf->ops->FetchAndOpBegin      = PetscSFFetchAndOpBegin_Basic;
629   sf->ops->FetchAndOpEnd        = PetscSFFetchAndOpEnd_Basic;
630   sf->ops->GetLeafRanks         = PetscSFGetLeafRanks_Basic;
631   sf->ops->CreateEmbeddedSF     = PetscSFCreateEmbeddedSF_Basic;
632   sf->ops->CreateEmbeddedLeafSF = PetscSFCreateEmbeddedLeafSF_Basic;
633 
634   ierr = PetscNewLog(sf,&dat);CHKERRQ(ierr);
635   sf->data = (void*)dat;
636   PetscFunctionReturn(0);
637 }
638