xref: /petsc/src/vec/is/sf/impls/basic/sfbasic.c (revision ffc4695bcb29f4b022f59a5fd6bc99fc280ff6d8)
18cd53115SBarry Smith 
220c24465SJunchao Zhang #include "petscsf.h"
340e23c03SJunchao Zhang #include <../src/vec/is/sf/impls/basic/sfbasic.h>
4cd620004SJunchao Zhang #include <../src/vec/is/sf/impls/basic/sfpack.h>
5b23bfdefSJunchao Zhang 
640e23c03SJunchao Zhang /*===================================================================================*/
740e23c03SJunchao Zhang /*              SF public interface implementations                                  */
840e23c03SJunchao Zhang /*===================================================================================*/
940e23c03SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFSetUp_Basic(PetscSF sf)
1095fce210SBarry Smith {
1195fce210SBarry Smith   PetscErrorCode ierr;
12b23bfdefSJunchao Zhang   PetscSF_Basic  *bas = (PetscSF_Basic*)sf->data;
1395fce210SBarry Smith   PetscInt       *rlengths,*ilengths,i;
1440e23c03SJunchao Zhang   PetscMPIInt    rank,niranks,*iranks,tag;
1595fce210SBarry Smith   MPI_Comm       comm;
16b5a8e515SJed Brown   MPI_Group      group;
1740e23c03SJunchao Zhang   MPI_Request    *rootreqs,*leafreqs;
1895fce210SBarry Smith 
1995fce210SBarry Smith   PetscFunctionBegin;
20*ffc4695bSBarry Smith   ierr = MPI_Comm_group(PETSC_COMM_SELF,&group);CHKERRMPI(ierr);
21b5a8e515SJed Brown   ierr = PetscSFSetUpRanks(sf,group);CHKERRQ(ierr);
22*ffc4695bSBarry Smith   ierr = MPI_Group_free(&group);CHKERRMPI(ierr);
2395fce210SBarry Smith   ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr);
2440e23c03SJunchao Zhang   ierr = PetscObjectGetNewTag((PetscObject)sf,&tag);CHKERRQ(ierr);
25*ffc4695bSBarry Smith   ierr = MPI_Comm_rank(comm,&rank);CHKERRMPI(ierr);
2695fce210SBarry Smith   /*
2795fce210SBarry Smith    * Inform roots about how many leaves and from which ranks
2895fce210SBarry Smith    */
29785e854fSJed Brown   ierr = PetscMalloc1(sf->nranks,&rlengths);CHKERRQ(ierr);
30cd620004SJunchao Zhang   /* Determine number, sending ranks and length of incoming */
3195fce210SBarry Smith   for (i=0; i<sf->nranks; i++) {
3295fce210SBarry Smith     rlengths[i] = sf->roffset[i+1] - sf->roffset[i]; /* Number of roots referenced by my leaves; for rank sf->ranks[i] */
3395fce210SBarry Smith   }
3440e23c03SJunchao Zhang   ierr = PetscCommBuildTwoSided(comm,1,MPIU_INT,sf->nranks-sf->ndranks,sf->ranks+sf->ndranks,rlengths+sf->ndranks,&niranks,&iranks,(void**)&ilengths);CHKERRQ(ierr);
35c943f53fSJed Brown 
360b899082SJunchao Zhang   /* Sort iranks. See use of VecScatterGetRemoteOrdered_Private() in MatGetBrowsOfAoCols_MPIAIJ() on why.
370b899082SJunchao Zhang      We could sort ranks there at the price of allocating extra working arrays. Presumably, niranks is
380b899082SJunchao Zhang      small and the sorting is cheap.
390b899082SJunchao Zhang    */
400b899082SJunchao Zhang   ierr = PetscSortMPIIntWithIntArray(niranks,iranks,ilengths);CHKERRQ(ierr);
410b899082SJunchao Zhang 
42c943f53fSJed Brown   /* Partition into distinguished and non-distinguished incoming ranks */
43c943f53fSJed Brown   bas->ndiranks = sf->ndranks;
44c943f53fSJed Brown   bas->niranks = bas->ndiranks + niranks;
45c943f53fSJed Brown   ierr = PetscMalloc2(bas->niranks,&bas->iranks,bas->niranks+1,&bas->ioffset);CHKERRQ(ierr);
46c943f53fSJed Brown   bas->ioffset[0] = 0;
47c943f53fSJed Brown   for (i=0; i<bas->ndiranks; i++) {
48c943f53fSJed Brown     bas->iranks[i] = sf->ranks[i];
49c943f53fSJed Brown     bas->ioffset[i+1] = bas->ioffset[i] + rlengths[i];
50c943f53fSJed Brown   }
5140e23c03SJunchao Zhang   if (bas->ndiranks > 1 || (bas->ndiranks == 1 && bas->iranks[0] != rank)) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Broken setup for shared ranks");
5240e23c03SJunchao Zhang   for (; i<bas->niranks; i++) {
53c943f53fSJed Brown     bas->iranks[i] = iranks[i-bas->ndiranks];
54c943f53fSJed Brown     bas->ioffset[i+1] = bas->ioffset[i] + ilengths[i-bas->ndiranks];
55c943f53fSJed Brown   }
56c943f53fSJed Brown   bas->itotal = bas->ioffset[i];
5795fce210SBarry Smith   ierr = PetscFree(rlengths);CHKERRQ(ierr);
58c943f53fSJed Brown   ierr = PetscFree(iranks);CHKERRQ(ierr);
59c943f53fSJed Brown   ierr = PetscFree(ilengths);CHKERRQ(ierr);
6095fce210SBarry Smith 
6195fce210SBarry Smith   /* Send leaf identities to roots */
62c943f53fSJed Brown   ierr = PetscMalloc1(bas->itotal,&bas->irootloc);CHKERRQ(ierr);
6340e23c03SJunchao Zhang   ierr = PetscMalloc2(bas->niranks-bas->ndiranks,&rootreqs,sf->nranks-sf->ndranks,&leafreqs);CHKERRQ(ierr);
6440e23c03SJunchao Zhang   for (i=bas->ndiranks; i<bas->niranks; i++) {
65*ffc4695bSBarry Smith     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]);CHKERRMPI(ierr);
6640e23c03SJunchao Zhang   }
6740e23c03SJunchao Zhang   for (i=0; i<sf->nranks; i++) {
6895fce210SBarry Smith     PetscMPIInt npoints;
6995fce210SBarry Smith     ierr = PetscMPIIntCast(sf->roffset[i+1] - sf->roffset[i],&npoints);CHKERRQ(ierr);
7040e23c03SJunchao Zhang     if (i < sf->ndranks) {
7140e23c03SJunchao Zhang       if (sf->ranks[i] != rank) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Cannot interpret distinguished leaf rank");
7240e23c03SJunchao Zhang       if (bas->iranks[0] != rank) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Cannot interpret distinguished root rank");
7340e23c03SJunchao Zhang       if (npoints != bas->ioffset[1]-bas->ioffset[0]) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Distinguished rank exchange has mismatched lengths");
7440e23c03SJunchao Zhang       ierr = PetscArraycpy(bas->irootloc+bas->ioffset[0],sf->rremote+sf->roffset[i],npoints);CHKERRQ(ierr);
75c943f53fSJed Brown       continue;
76c943f53fSJed Brown     }
77*ffc4695bSBarry Smith     ierr = MPI_Isend(sf->rremote+sf->roffset[i],npoints,MPIU_INT,sf->ranks[i],tag,comm,&leafreqs[i-sf->ndranks]);CHKERRMPI(ierr);
78bf39f1bfSJed Brown   }
79*ffc4695bSBarry Smith   ierr = MPI_Waitall(bas->niranks-bas->ndiranks,rootreqs,MPI_STATUSES_IGNORE);CHKERRMPI(ierr);
80*ffc4695bSBarry Smith   ierr = MPI_Waitall(sf->nranks-sf->ndranks,leafreqs,MPI_STATUSES_IGNORE);CHKERRMPI(ierr);
8140e23c03SJunchao Zhang   ierr = PetscFree2(rootreqs,leafreqs);CHKERRQ(ierr);
8295fce210SBarry Smith 
83cd620004SJunchao Zhang   sf->nleafreqs  = sf->nranks - sf->ndranks;
84cd620004SJunchao Zhang   bas->nrootreqs = bas->niranks - bas->ndiranks;
85cd620004SJunchao Zhang   sf->persistent = PETSC_TRUE;
86eb02082bSJunchao Zhang 
87cd620004SJunchao Zhang   /* Setup fields related to packing */
88cd620004SJunchao Zhang   ierr = PetscSFSetUpPackFields(sf);CHKERRQ(ierr);
8995fce210SBarry Smith   PetscFunctionReturn(0);
9095fce210SBarry Smith }
9195fce210SBarry Smith 
9240e23c03SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFReset_Basic(PetscSF sf)
9395fce210SBarry Smith {
9495fce210SBarry Smith   PetscErrorCode    ierr;
95cd620004SJunchao Zhang   PetscSF_Basic     *bas = (PetscSF_Basic*)sf->data;
9695fce210SBarry Smith 
9795fce210SBarry Smith   PetscFunctionBegin;
9829046d53SLisandro Dalcin   if (bas->inuse) SETERRQ(PetscObjectComm((PetscObject)sf),PETSC_ERR_ARG_WRONGSTATE,"Outstanding operation has not been completed");
99c943f53fSJed Brown   ierr = PetscFree2(bas->iranks,bas->ioffset);CHKERRQ(ierr);
100c943f53fSJed Brown   ierr = PetscFree(bas->irootloc);CHKERRQ(ierr);
1017fd2d3dbSJunchao Zhang #if defined(PETSC_HAVE_DEVICE)
10220c24465SJunchao Zhang   for (PetscInt i=0; i<2; i++) {ierr = PetscSFFree(sf,PETSC_MEMTYPE_DEVICE,bas->irootloc_d[i]);CHKERRQ(ierr);}
103eb02082bSJunchao Zhang #endif
104cd620004SJunchao Zhang   ierr = PetscSFLinkDestroy(sf,&bas->avail);CHKERRQ(ierr);
105cd620004SJunchao Zhang   ierr = PetscSFResetPackFields(sf);CHKERRQ(ierr);
10695fce210SBarry Smith   PetscFunctionReturn(0);
10795fce210SBarry Smith }
10895fce210SBarry Smith 
10940e23c03SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFDestroy_Basic(PetscSF sf)
11095fce210SBarry Smith {
11195fce210SBarry Smith   PetscErrorCode ierr;
11295fce210SBarry Smith 
11395fce210SBarry Smith   PetscFunctionBegin;
114f6d956f6SStefano Zampini   ierr = PetscSFReset_Basic(sf);CHKERRQ(ierr);
11595fce210SBarry Smith   ierr = PetscFree(sf->data);CHKERRQ(ierr);
11695fce210SBarry Smith   PetscFunctionReturn(0);
11795fce210SBarry Smith }
11895fce210SBarry Smith 
11940e23c03SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFView_Basic(PetscSF sf,PetscViewer viewer)
12095fce210SBarry Smith {
12195fce210SBarry Smith   PetscErrorCode ierr;
12295fce210SBarry Smith   PetscBool      iascii;
12395fce210SBarry Smith 
12495fce210SBarry Smith   PetscFunctionBegin;
12595fce210SBarry Smith   ierr = PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);CHKERRQ(ierr);
126b23bfdefSJunchao Zhang   if (iascii) {ierr = PetscViewerASCIIPrintf(viewer,"  sort=%s\n",sf->rankorder ? "rank-order" : "unordered");CHKERRQ(ierr);}
12795fce210SBarry Smith   PetscFunctionReturn(0);
12895fce210SBarry Smith }
12995fce210SBarry Smith 
130eb02082bSJunchao Zhang static PetscErrorCode PetscSFBcastAndOpBegin_Basic(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op)
13195fce210SBarry Smith {
13295fce210SBarry Smith   PetscErrorCode    ierr;
133cd620004SJunchao Zhang   PetscSFLink       link = NULL;
134851d6770SJunchao Zhang   MPI_Request       *rootreqs = NULL,*leafreqs = NULL;
135cd620004SJunchao Zhang   PetscSF_Basic     *bas = (PetscSF_Basic*)sf->data;
13695fce210SBarry Smith 
13795fce210SBarry Smith   PetscFunctionBegin;
138cd620004SJunchao Zhang   /* Create a communication link, which provides buffers & MPI requests etc */
139cd620004SJunchao Zhang   ierr = PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,PETSCSF_BCAST,&link);CHKERRQ(ierr);
140cd620004SJunchao Zhang   /* Get MPI requests from the link. We do not need buffers explicitly since we use persistent MPI */
141cd620004SJunchao Zhang   ierr = PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_ROOT2LEAF,NULL,NULL,&rootreqs,&leafreqs);CHKERRQ(ierr);
142cd620004SJunchao Zhang   /* Post Irecv for remote */
143*ffc4695bSBarry Smith   ierr = MPI_Startall_irecv(sf->leafbuflen[PETSCSF_REMOTE],unit,sf->nleafreqs,leafreqs);CHKERRMPI(ierr);
144cd620004SJunchao Zhang   /* Pack rootdata and do Isend for remote */
145cd620004SJunchao Zhang   ierr = PetscSFLinkPackRootData(sf,link,PETSCSF_REMOTE,rootdata);CHKERRQ(ierr);
146*ffc4695bSBarry Smith   ierr = MPI_Startall_isend(bas->rootbuflen[PETSCSF_REMOTE],unit,bas->nrootreqs,rootreqs);CHKERRMPI(ierr);
147cd620004SJunchao Zhang   /* Do local BcastAndOp, which overlaps with the irecv/isend above */
148cd620004SJunchao Zhang   ierr = PetscSFLinkBcastAndOpLocal(sf,link,rootdata,leafdata,op);CHKERRQ(ierr);
14995fce210SBarry Smith   PetscFunctionReturn(0);
15095fce210SBarry Smith }
15195fce210SBarry Smith 
15200816365SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFBcastAndOpEnd_Basic(PetscSF sf,MPI_Datatype unit,const void *rootdata,void *leafdata,MPI_Op op)
15395fce210SBarry Smith {
15495fce210SBarry Smith   PetscErrorCode    ierr;
155cd620004SJunchao Zhang   PetscSFLink       link = NULL;
15695fce210SBarry Smith 
15795fce210SBarry Smith   PetscFunctionBegin;
158cd620004SJunchao Zhang   /* Retrieve the link used in XxxBegin() with root/leafdata as key */
159cd620004SJunchao Zhang   ierr = PetscSFLinkGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
160cd620004SJunchao Zhang   /* Wait for the completion of mpi */
161cd620004SJunchao Zhang   ierr = PetscSFLinkMPIWaitall(sf,link,PETSCSF_ROOT2LEAF);CHKERRQ(ierr);
162cd620004SJunchao Zhang   /* Unpack leafdata and reclaim the link */
163cd620004SJunchao Zhang   ierr = PetscSFLinkUnpackLeafData(sf,link,PETSCSF_REMOTE,leafdata,op);CHKERRQ(ierr);
164cd620004SJunchao Zhang   ierr = PetscSFLinkReclaim(sf,&link);CHKERRQ(ierr);
165cd620004SJunchao Zhang   PetscFunctionReturn(0);
166cd620004SJunchao Zhang }
167cd620004SJunchao Zhang 
168cd620004SJunchao Zhang /* Shared by ReduceBegin and FetchAndOpBegin */
169cd620004SJunchao Zhang PETSC_STATIC_INLINE PetscErrorCode PetscSFLeafToRootBegin_Basic(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op,PetscSFOperation sfop,PetscSFLink *out)
170cd620004SJunchao Zhang {
171cd620004SJunchao Zhang   PetscErrorCode    ierr;
172cd620004SJunchao Zhang   PetscSFLink       link;
173cd620004SJunchao Zhang   PetscSF_Basic     *bas = (PetscSF_Basic*)sf->data;
174cd620004SJunchao Zhang   MPI_Request       *rootreqs = NULL,*leafreqs = NULL;
175cd620004SJunchao Zhang 
176cd620004SJunchao Zhang   PetscFunctionBegin;
177cd620004SJunchao Zhang   ierr = PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,sfop,&link);CHKERRQ(ierr);
178cd620004SJunchao Zhang   ierr = PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_LEAF2ROOT,NULL,NULL,&rootreqs,&leafreqs);CHKERRQ(ierr);
179*ffc4695bSBarry Smith   ierr = MPI_Startall_irecv(bas->rootbuflen[PETSCSF_REMOTE],unit,bas->nrootreqs,rootreqs);CHKERRMPI(ierr);
180cd620004SJunchao Zhang   ierr = PetscSFLinkPackLeafData(sf,link,PETSCSF_REMOTE,leafdata);CHKERRQ(ierr);
181*ffc4695bSBarry Smith   ierr = MPI_Startall_isend(sf->leafbuflen[PETSCSF_REMOTE],unit,sf->nleafreqs,leafreqs);CHKERRMPI(ierr);
182cd620004SJunchao Zhang   *out = link;
18395fce210SBarry Smith   PetscFunctionReturn(0);
18495fce210SBarry Smith }
18595fce210SBarry Smith 
18695fce210SBarry Smith /* leaf -> root with reduction */
187eb02082bSJunchao Zhang static PetscErrorCode PetscSFReduceBegin_Basic(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op)
18895fce210SBarry Smith {
18995fce210SBarry Smith   PetscErrorCode    ierr;
190cd620004SJunchao Zhang   PetscSFLink       link = NULL;
19195fce210SBarry Smith 
19295fce210SBarry Smith   PetscFunctionBegin;
193cd620004SJunchao Zhang   ierr = PetscSFLeafToRootBegin_Basic(sf,unit,leafmtype,leafdata,rootmtype,rootdata,op,PETSCSF_REDUCE,&link);CHKERRQ(ierr);
194cd620004SJunchao Zhang   ierr = PetscSFLinkReduceLocal(sf,link,leafdata,rootdata,op);CHKERRQ(ierr);
19595fce210SBarry Smith   PetscFunctionReturn(0);
19695fce210SBarry Smith }
19795fce210SBarry Smith 
19800816365SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFReduceEnd_Basic(PetscSF sf,MPI_Datatype unit,const void *leafdata,void *rootdata,MPI_Op op)
19995fce210SBarry Smith {
20095fce210SBarry Smith   PetscErrorCode    ierr;
201cd620004SJunchao Zhang   PetscSFLink       link = NULL;
20295fce210SBarry Smith 
20395fce210SBarry Smith   PetscFunctionBegin;
204cd620004SJunchao Zhang   ierr = PetscSFLinkGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
205cd620004SJunchao Zhang   ierr = PetscSFLinkMPIWaitall(sf,link,PETSCSF_LEAF2ROOT);CHKERRQ(ierr);
206cd620004SJunchao Zhang   ierr = PetscSFLinkUnpackRootData(sf,link,PETSCSF_REMOTE,rootdata,op);CHKERRQ(ierr);
207cd620004SJunchao Zhang   ierr = PetscSFLinkReclaim(sf,&link);CHKERRQ(ierr);
20895fce210SBarry Smith   PetscFunctionReturn(0);
20995fce210SBarry Smith }
21095fce210SBarry Smith 
211eb02082bSJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFFetchAndOpBegin_Basic(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,void *rootdata,PetscMemType leafmtype,const void *leafdata,void *leafupdate,MPI_Op op)
21295fce210SBarry Smith {
21395fce210SBarry Smith   PetscErrorCode    ierr;
214cd620004SJunchao Zhang   PetscSFLink       link = NULL;
21595fce210SBarry Smith 
21695fce210SBarry Smith   PetscFunctionBegin;
217cd620004SJunchao Zhang   ierr = PetscSFLeafToRootBegin_Basic(sf,unit,leafmtype,leafdata,rootmtype,rootdata,op,PETSCSF_FETCH,&link);CHKERRQ(ierr);
218cd620004SJunchao Zhang   ierr = PetscSFLinkFetchAndOpLocal(sf,link,rootdata,leafdata,leafupdate,op);CHKERRQ(ierr);
21995fce210SBarry Smith   PetscFunctionReturn(0);
22095fce210SBarry Smith }
22195fce210SBarry Smith 
22200816365SJunchao Zhang static PetscErrorCode PetscSFFetchAndOpEnd_Basic(PetscSF sf,MPI_Datatype unit,void *rootdata,const void *leafdata,void *leafupdate,MPI_Op op)
22395fce210SBarry Smith {
22495fce210SBarry Smith   PetscErrorCode    ierr;
225cd620004SJunchao Zhang   PetscSFLink       link = NULL;
226851d6770SJunchao Zhang   MPI_Request       *rootreqs = NULL,*leafreqs = NULL;
227cd620004SJunchao Zhang   PetscSF_Basic     *bas = (PetscSF_Basic*)sf->data;
22895fce210SBarry Smith 
22995fce210SBarry Smith   PetscFunctionBegin;
230cd620004SJunchao Zhang   ierr = PetscSFLinkGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
23195fce210SBarry Smith   /* This implementation could be changed to unpack as receives arrive, at the cost of non-determinism */
232cd620004SJunchao Zhang   ierr = PetscSFLinkMPIWaitall(sf,link,PETSCSF_LEAF2ROOT);CHKERRQ(ierr);
233cd620004SJunchao Zhang   /* Do fetch-and-op, the (remote) update results are in rootbuf */
234cd620004SJunchao Zhang   ierr = PetscSFLinkFetchRootData(sf,link,PETSCSF_REMOTE,rootdata,op);CHKERRQ(ierr);
23540e23c03SJunchao Zhang 
236cd620004SJunchao Zhang   /* Bcast rootbuf to leafupdate */
237cd620004SJunchao Zhang   ierr = PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_ROOT2LEAF,NULL,NULL,&rootreqs,&leafreqs);CHKERRQ(ierr);
238cd620004SJunchao Zhang   /* Post leaf receives and root sends */
239*ffc4695bSBarry Smith   ierr = MPI_Startall_irecv(sf->leafbuflen[PETSCSF_REMOTE],unit,sf->nleafreqs,leafreqs);CHKERRMPI(ierr);
240*ffc4695bSBarry Smith   ierr = MPI_Startall_isend(bas->rootbuflen[PETSCSF_REMOTE],unit,bas->nrootreqs,rootreqs);CHKERRMPI(ierr);
241b23bfdefSJunchao Zhang   /* Unpack and insert fetched data into leaves */
242cd620004SJunchao Zhang   ierr = PetscSFLinkMPIWaitall(sf,link,PETSCSF_ROOT2LEAF);CHKERRQ(ierr);
243cd620004SJunchao Zhang   ierr = PetscSFLinkUnpackLeafData(sf,link,PETSCSF_REMOTE,leafupdate,MPIU_REPLACE);CHKERRQ(ierr);
244cd620004SJunchao Zhang   ierr = PetscSFLinkReclaim(sf,&link);CHKERRQ(ierr);
24595fce210SBarry Smith   PetscFunctionReturn(0);
24695fce210SBarry Smith }
24795fce210SBarry Smith 
24840e23c03SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFGetLeafRanks_Basic(PetscSF sf,PetscInt *niranks,const PetscMPIInt **iranks,const PetscInt **ioffset,const PetscInt **irootloc)
2498750ddebSJunchao Zhang {
2508750ddebSJunchao Zhang   PetscSF_Basic *bas = (PetscSF_Basic*)sf->data;
2518750ddebSJunchao Zhang 
2528750ddebSJunchao Zhang   PetscFunctionBegin;
2538750ddebSJunchao Zhang   if (niranks)  *niranks  = bas->niranks;
2548750ddebSJunchao Zhang   if (iranks)   *iranks   = bas->iranks;
2558750ddebSJunchao Zhang   if (ioffset)  *ioffset  = bas->ioffset;
2568750ddebSJunchao Zhang   if (irootloc) *irootloc = bas->irootloc;
2578750ddebSJunchao Zhang   PetscFunctionReturn(0);
2588750ddebSJunchao Zhang }
2598750ddebSJunchao Zhang 
260f659e5c7SJunchao Zhang /* An optimized PetscSFCreateEmbeddedSF. We aggresively make use of the established communication on sf.
261f659e5c7SJunchao Zhang    We need one bcast on sf, and no communication anymore to build the embedded sf. Note that selected[]
262f659e5c7SJunchao Zhang    was sorted before calling the routine.
263f659e5c7SJunchao Zhang  */
264f659e5c7SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFCreateEmbeddedSF_Basic(PetscSF sf,PetscInt nselected,const PetscInt *selected,PetscSF *newsf)
265f659e5c7SJunchao Zhang {
266f659e5c7SJunchao Zhang   PetscSF           esf;
267cd620004SJunchao Zhang   PetscInt          esf_nranks,esf_ndranks,*esf_roffset,*esf_rmine,*esf_rremote;
268cd620004SJunchao Zhang   PetscInt          i,j,p,q,nroots,esf_nleaves,*new_ilocal,nranks,ndranks,niranks,ndiranks,minleaf,maxleaf,maxlocal;
269cd620004SJunchao Zhang   char              *rootdata,*leafdata,*leafmem; /* Only stores 0 or 1, so we can save memory with char */
270f659e5c7SJunchao Zhang   PetscMPIInt       *esf_ranks;
271f659e5c7SJunchao Zhang   const PetscMPIInt *ranks,*iranks;
272cd620004SJunchao Zhang   const PetscInt    *roffset,*rmine,*rremote,*ioffset,*irootloc;
273f659e5c7SJunchao Zhang   PetscBool         connected;
274f659e5c7SJunchao Zhang   PetscSFNode       *new_iremote;
275f659e5c7SJunchao Zhang   PetscSF_Basic     *bas;
276f659e5c7SJunchao Zhang   PetscErrorCode    ierr;
277f659e5c7SJunchao Zhang 
278f659e5c7SJunchao Zhang   PetscFunctionBegin;
279f659e5c7SJunchao Zhang   ierr = PetscSFCreate(PetscObjectComm((PetscObject)sf),&esf);CHKERRQ(ierr);
28020c24465SJunchao Zhang   ierr = PetscSFSetFromOptions(esf);CHKERRQ(ierr);
281f659e5c7SJunchao Zhang   ierr = PetscSFSetType(esf,PETSCSFBASIC);CHKERRQ(ierr); /* This optimized routine can only create a basic sf */
282f659e5c7SJunchao Zhang 
283cd620004SJunchao Zhang   /* Find out which leaves are still connected to roots in the embedded sf by doing a Bcast */
284f659e5c7SJunchao Zhang   ierr = PetscSFGetGraph(sf,&nroots,NULL,NULL,NULL);CHKERRQ(ierr);
285f659e5c7SJunchao Zhang   ierr = PetscSFGetLeafRange(sf,&minleaf,&maxleaf);CHKERRQ(ierr);
286cd620004SJunchao Zhang   maxlocal = maxleaf - minleaf + 1;
287cd620004SJunchao Zhang   ierr = PetscCalloc2(nroots,&rootdata,maxlocal,&leafmem);CHKERRQ(ierr);
288cd620004SJunchao Zhang   leafdata = leafmem - minleaf;
289f659e5c7SJunchao Zhang   /* Tag selected roots */
290f659e5c7SJunchao Zhang   for (i=0; i<nselected; ++i) rootdata[selected[i]] = 1;
291f659e5c7SJunchao Zhang 
292cd620004SJunchao Zhang   ierr = PetscSFBcastBegin(sf,MPI_CHAR,rootdata,leafdata);CHKERRQ(ierr);
293cd620004SJunchao Zhang   ierr = PetscSFBcastEnd(sf,MPI_CHAR,rootdata,leafdata);CHKERRQ(ierr);
294f659e5c7SJunchao Zhang   ierr = PetscSFGetLeafInfo_Basic(sf,&nranks,&ndranks,&ranks,&roffset,&rmine,&rremote);CHKERRQ(ierr); /* Get send info */
295cd620004SJunchao Zhang   esf_nranks = esf_ndranks = esf_nleaves = 0;
296b23bfdefSJunchao Zhang   for (i=0; i<nranks; i++) {
297cd620004SJunchao Zhang     connected = PETSC_FALSE; /* Is this process still connected to this remote root rank? */
298cd620004SJunchao Zhang     for (j=roffset[i]; j<roffset[i+1]; j++) {if (leafdata[rmine[j]]) {esf_nleaves++; connected = PETSC_TRUE;}}
299f659e5c7SJunchao Zhang     if (connected) {esf_nranks++; if (i < ndranks) esf_ndranks++;}
300f659e5c7SJunchao Zhang   }
301f659e5c7SJunchao Zhang 
302f659e5c7SJunchao Zhang   /* Set graph of esf and also set up its outgoing communication (i.e., send info), which is usually done by PetscSFSetUpRanks */
303cd620004SJunchao Zhang   ierr = PetscMalloc1(esf_nleaves,&new_ilocal);CHKERRQ(ierr);
304cd620004SJunchao Zhang   ierr = PetscMalloc1(esf_nleaves,&new_iremote);CHKERRQ(ierr);
305cd620004SJunchao Zhang   ierr = PetscMalloc4(esf_nranks,&esf_ranks,esf_nranks+1,&esf_roffset,esf_nleaves,&esf_rmine,esf_nleaves,&esf_rremote);CHKERRQ(ierr);
306f659e5c7SJunchao Zhang   p    = 0; /* Counter for connected root ranks */
307f659e5c7SJunchao Zhang   q    = 0; /* Counter for connected leaves */
308f659e5c7SJunchao Zhang   esf_roffset[0] = 0;
309f659e5c7SJunchao Zhang   for (i=0; i<nranks; i++) { /* Scan leaf data again to fill esf arrays */
310f659e5c7SJunchao Zhang     connected = PETSC_FALSE;
311cd620004SJunchao Zhang     for (j=roffset[i]; j<roffset[i+1]; j++) {
312cd620004SJunchao Zhang       if (leafdata[rmine[j]]) {
313f659e5c7SJunchao Zhang         esf_rmine[q]         = new_ilocal[q] = rmine[j];
314f659e5c7SJunchao Zhang         esf_rremote[q]       = rremote[j];
315f659e5c7SJunchao Zhang         new_iremote[q].index = rremote[j];
316f659e5c7SJunchao Zhang         new_iremote[q].rank  = ranks[i];
317f659e5c7SJunchao Zhang         connected            = PETSC_TRUE;
318f659e5c7SJunchao Zhang         q++;
319f659e5c7SJunchao Zhang       }
320f659e5c7SJunchao Zhang     }
321f659e5c7SJunchao Zhang     if (connected) {
322f659e5c7SJunchao Zhang       esf_ranks[p]     = ranks[i];
323f659e5c7SJunchao Zhang       esf_roffset[p+1] = q;
324f659e5c7SJunchao Zhang       p++;
325f659e5c7SJunchao Zhang     }
326f659e5c7SJunchao Zhang   }
327f659e5c7SJunchao Zhang 
328f659e5c7SJunchao Zhang   /* SetGraph internally resets the SF, so we only set its fields after the call */
329cd620004SJunchao Zhang   ierr           = PetscSFSetGraph(esf,nroots,esf_nleaves,new_ilocal,PETSC_OWN_POINTER,new_iremote,PETSC_OWN_POINTER);CHKERRQ(ierr);
330f659e5c7SJunchao Zhang   esf->nranks    = esf_nranks;
331f659e5c7SJunchao Zhang   esf->ndranks   = esf_ndranks;
332f659e5c7SJunchao Zhang   esf->ranks     = esf_ranks;
333f659e5c7SJunchao Zhang   esf->roffset   = esf_roffset;
334f659e5c7SJunchao Zhang   esf->rmine     = esf_rmine;
335f659e5c7SJunchao Zhang   esf->rremote   = esf_rremote;
336cd620004SJunchao Zhang   esf->nleafreqs = esf_nranks - esf_ndranks;
337f659e5c7SJunchao Zhang 
338f659e5c7SJunchao Zhang   /* Set up the incoming communication (i.e., recv info) stored in esf->data, which is usually done by PetscSFSetUp_Basic */
339f659e5c7SJunchao Zhang   bas  = (PetscSF_Basic*)esf->data;
340f659e5c7SJunchao Zhang   ierr = PetscSFGetRootInfo_Basic(sf,&niranks,&ndiranks,&iranks,&ioffset,&irootloc);CHKERRQ(ierr); /* Get recv info */
341f659e5c7SJunchao Zhang   /* Embedded sf always has simpler communication than the original one. We might allocate longer arrays than needed here. But we
342cd620004SJunchao Zhang      we do not care since these arrays are usually short. The benefit is we can fill these arrays by just parsing irootloc once.
343f659e5c7SJunchao Zhang    */
344f659e5c7SJunchao Zhang   ierr = PetscMalloc2(niranks,&bas->iranks,niranks+1,&bas->ioffset);CHKERRQ(ierr);
345f659e5c7SJunchao Zhang   ierr = PetscMalloc1(ioffset[niranks],&bas->irootloc);CHKERRQ(ierr);
346f659e5c7SJunchao Zhang   bas->niranks = bas->ndiranks = bas->ioffset[0] = 0;
347f659e5c7SJunchao Zhang   p = 0; /* Counter for connected leaf ranks */
348f659e5c7SJunchao Zhang   q = 0; /* Counter for connected roots */
349f659e5c7SJunchao Zhang   for (i=0; i<niranks; i++) {
350f659e5c7SJunchao Zhang     connected = PETSC_FALSE; /* Is the current process still connected to this remote leaf rank? */
351f659e5c7SJunchao Zhang     for (j=ioffset[i]; j<ioffset[i+1]; j++) {
352cd620004SJunchao Zhang       if (rootdata[irootloc[j]]) {
353f659e5c7SJunchao Zhang         bas->irootloc[q++] = irootloc[j];
354f659e5c7SJunchao Zhang         connected = PETSC_TRUE;
355f659e5c7SJunchao Zhang       }
356f659e5c7SJunchao Zhang     }
357f659e5c7SJunchao Zhang     if (connected) {
358f659e5c7SJunchao Zhang       bas->niranks++;
359f659e5c7SJunchao Zhang       if (i<ndiranks) bas->ndiranks++; /* Note that order of ranks (including distinguished ranks) is kept */
360f659e5c7SJunchao Zhang       bas->iranks[p]    = iranks[i];
361f659e5c7SJunchao Zhang       bas->ioffset[p+1] = q;
362f659e5c7SJunchao Zhang       p++;
363f659e5c7SJunchao Zhang     }
364f659e5c7SJunchao Zhang   }
365f659e5c7SJunchao Zhang   bas->itotal     = q;
366cd620004SJunchao Zhang   bas->nrootreqs  = bas->niranks - bas->ndiranks;
367cd620004SJunchao Zhang   esf->persistent = PETSC_TRUE;
368cd620004SJunchao Zhang   /* Setup packing related fields */
369cd620004SJunchao Zhang   ierr = PetscSFSetUpPackFields(esf);CHKERRQ(ierr);
370f659e5c7SJunchao Zhang 
37120c24465SJunchao Zhang   /* Copy from PetscSFSetUp(), since this method wants to skip PetscSFSetUp(). */
37220c24465SJunchao Zhang #if defined(PETSC_HAVE_CUDA)
37320c24465SJunchao Zhang   if (esf->backend == PETSCSF_BACKEND_CUDA) {
37420c24465SJunchao Zhang     esf->ops->Malloc = PetscSFMalloc_Cuda;
37520c24465SJunchao Zhang     esf->ops->Free   = PetscSFFree_Cuda;
37620c24465SJunchao Zhang   }
37720c24465SJunchao Zhang #endif
37820c24465SJunchao Zhang 
37959af0bd3SScott Kruger #if defined(PETSC_HAVE_HIP)
38059af0bd3SScott Kruger   /* TODO: Needs debugging */
38159af0bd3SScott Kruger   if (esf->backend == PETSCSF_BACKEND_HIP) {
38259af0bd3SScott Kruger     esf->ops->Malloc = PetscSFMalloc_HIP;
38359af0bd3SScott Kruger     esf->ops->Free   = PetscSFFree_HIP;
38459af0bd3SScott Kruger   }
38559af0bd3SScott Kruger #endif
38659af0bd3SScott Kruger 
38720c24465SJunchao Zhang #if defined(PETSC_HAVE_KOKKOS)
38820c24465SJunchao Zhang   if (esf->backend == PETSCSF_BACKEND_KOKKOS) {
38920c24465SJunchao Zhang     esf->ops->Malloc = PetscSFMalloc_Kokkos;
39020c24465SJunchao Zhang     esf->ops->Free   = PetscSFFree_Kokkos;
39120c24465SJunchao Zhang   }
39220c24465SJunchao Zhang #endif
393f659e5c7SJunchao Zhang   esf->setupcalled = PETSC_TRUE; /* We have done setup ourselves! */
394cd620004SJunchao Zhang   ierr = PetscFree2(rootdata,leafmem);CHKERRQ(ierr);
395f659e5c7SJunchao Zhang   *newsf = esf;
396f659e5c7SJunchao Zhang   PetscFunctionReturn(0);
397f659e5c7SJunchao Zhang }
398f659e5c7SJunchao Zhang 
3998cc058d9SJed Brown PETSC_EXTERN PetscErrorCode PetscSFCreate_Basic(PetscSF sf)
40095fce210SBarry Smith {
40140e23c03SJunchao Zhang   PetscSF_Basic  *dat;
40295fce210SBarry Smith   PetscErrorCode ierr;
40395fce210SBarry Smith 
40495fce210SBarry Smith   PetscFunctionBegin;
40595fce210SBarry Smith   sf->ops->SetUp                = PetscSFSetUp_Basic;
40695fce210SBarry Smith   sf->ops->Reset                = PetscSFReset_Basic;
40795fce210SBarry Smith   sf->ops->Destroy              = PetscSFDestroy_Basic;
40895fce210SBarry Smith   sf->ops->View                 = PetscSFView_Basic;
4093482bfa8SJunchao Zhang   sf->ops->BcastAndOpBegin      = PetscSFBcastAndOpBegin_Basic;
4103482bfa8SJunchao Zhang   sf->ops->BcastAndOpEnd        = PetscSFBcastAndOpEnd_Basic;
41195fce210SBarry Smith   sf->ops->ReduceBegin          = PetscSFReduceBegin_Basic;
41295fce210SBarry Smith   sf->ops->ReduceEnd            = PetscSFReduceEnd_Basic;
41395fce210SBarry Smith   sf->ops->FetchAndOpBegin      = PetscSFFetchAndOpBegin_Basic;
41495fce210SBarry Smith   sf->ops->FetchAndOpEnd        = PetscSFFetchAndOpEnd_Basic;
4158750ddebSJunchao Zhang   sf->ops->GetLeafRanks         = PetscSFGetLeafRanks_Basic;
416f659e5c7SJunchao Zhang   sf->ops->CreateEmbeddedSF     = PetscSFCreateEmbeddedSF_Basic;
41795fce210SBarry Smith 
41840e23c03SJunchao Zhang   ierr = PetscNewLog(sf,&dat);CHKERRQ(ierr);
41940e23c03SJunchao Zhang   sf->data = (void*)dat;
42095fce210SBarry Smith   PetscFunctionReturn(0);
42195fce210SBarry Smith }
422