xref: /petsc/src/vec/is/sf/impls/basic/allgatherv/sfallgatherv.c (revision 2c71b3e237ead271e4f3aa1505f92bf476e3413d)
1dd5b3ca6SJunchao Zhang #include <../src/vec/is/sf/impls/basic/allgatherv/sfallgatherv.h>
2dd5b3ca6SJunchao Zhang 
3ad227feaSJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFBcastBegin_Gatherv(PetscSF,MPI_Datatype,PetscMemType,const void*,PetscMemType,void*,MPI_Op);
4dd5b3ca6SJunchao Zhang 
5dd5b3ca6SJunchao Zhang /* PetscSFGetGraph is non-collective. An implementation should not have collective calls */
6dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFGetGraph_Allgatherv(PetscSF sf,PetscInt *nroots,PetscInt *nleaves,const PetscInt **ilocal,const PetscSFNode **iremote)
7dd5b3ca6SJunchao Zhang {
8dd5b3ca6SJunchao Zhang   PetscErrorCode ierr;
9dd5b3ca6SJunchao Zhang   PetscInt       i,j,k;
10dd5b3ca6SJunchao Zhang   const PetscInt *range;
11dd5b3ca6SJunchao Zhang   PetscMPIInt    size;
12dd5b3ca6SJunchao Zhang 
13dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
14ffc4695bSBarry Smith   ierr = MPI_Comm_size(PetscObjectComm((PetscObject)sf),&size);CHKERRMPI(ierr);
15dd5b3ca6SJunchao Zhang   if (nroots)  *nroots  = sf->nroots;
16dd5b3ca6SJunchao Zhang   if (nleaves) *nleaves = sf->nleaves;
17dd5b3ca6SJunchao Zhang   if (ilocal)  *ilocal  = NULL; /* Contiguous leaves */
18dd5b3ca6SJunchao Zhang   if (iremote) {
19dd5b3ca6SJunchao Zhang     if (!sf->remote && sf->nleaves) { /* The && sf->nleaves makes sfgatherv able to inherit this routine */
20dd5b3ca6SJunchao Zhang       ierr = PetscLayoutGetRanges(sf->map,&range);CHKERRQ(ierr);
21dd5b3ca6SJunchao Zhang       ierr = PetscMalloc1(sf->nleaves,&sf->remote);CHKERRQ(ierr);
22dd5b3ca6SJunchao Zhang       sf->remote_alloc = sf->remote;
23dd5b3ca6SJunchao Zhang       for (i=0; i<size; i++) {
24dd5b3ca6SJunchao Zhang         for (j=range[i],k=0; j<range[i+1]; j++,k++) {
25dd5b3ca6SJunchao Zhang           sf->remote[j].rank  = i;
26dd5b3ca6SJunchao Zhang           sf->remote[j].index = k;
27dd5b3ca6SJunchao Zhang         }
28dd5b3ca6SJunchao Zhang       }
29dd5b3ca6SJunchao Zhang     }
30dd5b3ca6SJunchao Zhang     *iremote = sf->remote;
31dd5b3ca6SJunchao Zhang   }
32dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
33dd5b3ca6SJunchao Zhang }
34dd5b3ca6SJunchao Zhang 
35dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFSetUp_Allgatherv(PetscSF sf)
36dd5b3ca6SJunchao Zhang {
37dd5b3ca6SJunchao Zhang   PetscErrorCode     ierr;
38dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv*)sf->data;
39dd5b3ca6SJunchao Zhang   PetscMPIInt        size;
40dd5b3ca6SJunchao Zhang   PetscInt           i;
41dd5b3ca6SJunchao Zhang   const PetscInt     *range;
42dd5b3ca6SJunchao Zhang 
43dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
44cd620004SJunchao Zhang   ierr = PetscSFSetUp_Allgather(sf);CHKERRQ(ierr);
45ffc4695bSBarry Smith   ierr = MPI_Comm_size(PetscObjectComm((PetscObject)sf),&size);CHKERRMPI(ierr);
46dd5b3ca6SJunchao Zhang   if (sf->nleaves) { /* This if (sf->nleaves) test makes sfgatherv able to inherit this routine */
47dd5b3ca6SJunchao Zhang     ierr = PetscMalloc1(size,&dat->recvcounts);CHKERRQ(ierr);
48dd5b3ca6SJunchao Zhang     ierr = PetscMalloc1(size,&dat->displs);CHKERRQ(ierr);
49dd5b3ca6SJunchao Zhang     ierr = PetscLayoutGetRanges(sf->map,&range);CHKERRQ(ierr);
50dd5b3ca6SJunchao Zhang 
51dd5b3ca6SJunchao Zhang     for (i=0; i<size; i++) {
52dd5b3ca6SJunchao Zhang       ierr = PetscMPIIntCast(range[i],&dat->displs[i]);CHKERRQ(ierr);
53dd5b3ca6SJunchao Zhang       ierr = PetscMPIIntCast(range[i+1]-range[i],&dat->recvcounts[i]);CHKERRQ(ierr);
54dd5b3ca6SJunchao Zhang     }
55dd5b3ca6SJunchao Zhang   }
56dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
57dd5b3ca6SJunchao Zhang }
58dd5b3ca6SJunchao Zhang 
59dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFReset_Allgatherv(PetscSF sf)
60dd5b3ca6SJunchao Zhang {
61dd5b3ca6SJunchao Zhang   PetscErrorCode         ierr;
62eb02082bSJunchao Zhang   PetscSF_Allgatherv     *dat = (PetscSF_Allgatherv*)sf->data;
6371438e86SJunchao Zhang   PetscSFLink            link = dat->avail,next;
64dd5b3ca6SJunchao Zhang 
65dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
66dd5b3ca6SJunchao Zhang   ierr = PetscFree(dat->iranks);CHKERRQ(ierr);
67dd5b3ca6SJunchao Zhang   ierr = PetscFree(dat->ioffset);CHKERRQ(ierr);
68dd5b3ca6SJunchao Zhang   ierr = PetscFree(dat->irootloc);CHKERRQ(ierr);
69dd5b3ca6SJunchao Zhang   ierr = PetscFree(dat->recvcounts);CHKERRQ(ierr);
70dd5b3ca6SJunchao Zhang   ierr = PetscFree(dat->displs);CHKERRQ(ierr);
71*2c71b3e2SJacob Faibussowitsch   PetscCheckFalse(dat->inuse,PetscObjectComm((PetscObject)sf),PETSC_ERR_ARG_WRONGSTATE,"Outstanding operation has not been completed");
7271438e86SJunchao Zhang   for (; link; link=next) {next = link->next; ierr = PetscSFLinkDestroy(sf,link);CHKERRQ(ierr);}
7371438e86SJunchao Zhang   dat->avail = NULL;
74dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
75dd5b3ca6SJunchao Zhang }
76dd5b3ca6SJunchao Zhang 
77dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFDestroy_Allgatherv(PetscSF sf)
78dd5b3ca6SJunchao Zhang {
79dd5b3ca6SJunchao Zhang   PetscErrorCode ierr;
80dd5b3ca6SJunchao Zhang 
81dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
82dd5b3ca6SJunchao Zhang   ierr = PetscSFReset_Allgatherv(sf);CHKERRQ(ierr);
83dd5b3ca6SJunchao Zhang   ierr = PetscFree(sf->data);CHKERRQ(ierr);
84dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
85dd5b3ca6SJunchao Zhang }
86dd5b3ca6SJunchao Zhang 
87ad227feaSJunchao Zhang static PetscErrorCode PetscSFBcastBegin_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op)
88dd5b3ca6SJunchao Zhang {
89dd5b3ca6SJunchao Zhang   PetscErrorCode         ierr;
90cd620004SJunchao Zhang   PetscSFLink            link;
91dd5b3ca6SJunchao Zhang   PetscMPIInt            sendcount;
92dd5b3ca6SJunchao Zhang   MPI_Comm               comm;
93cd620004SJunchao Zhang   void                   *rootbuf = NULL,*leafbuf = NULL;
94cd620004SJunchao Zhang   MPI_Request            *req;
95dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv     *dat = (PetscSF_Allgatherv*)sf->data;
96dd5b3ca6SJunchao Zhang 
97dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
98cd620004SJunchao Zhang   ierr = PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,PETSCSF_BCAST,&link);CHKERRQ(ierr);
99cd620004SJunchao Zhang   ierr = PetscSFLinkPackRootData(sf,link,PETSCSF_REMOTE,rootdata);CHKERRQ(ierr);
10071438e86SJunchao Zhang   ierr = PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_TRUE/* device2host before sending */);CHKERRQ(ierr);
101dd5b3ca6SJunchao Zhang   ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr);
102dd5b3ca6SJunchao Zhang   ierr = PetscMPIIntCast(sf->nroots,&sendcount);CHKERRQ(ierr);
103cd620004SJunchao Zhang   ierr = PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_ROOT2LEAF,&rootbuf,&leafbuf,&req,NULL);CHKERRQ(ierr);
10471438e86SJunchao Zhang   ierr = PetscSFLinkSyncStreamBeforeCallMPI(sf,link,PETSCSF_ROOT2LEAF);CHKERRQ(ierr);
105820f2d46SBarry Smith   ierr = MPIU_Iallgatherv(rootbuf,sendcount,unit,leafbuf,dat->recvcounts,dat->displs,unit,comm,req);CHKERRMPI(ierr);
106855db38dSJunchao Zhang   PetscFunctionReturn(0);
107855db38dSJunchao Zhang }
108855db38dSJunchao Zhang 
109eb02082bSJunchao Zhang static PetscErrorCode PetscSFReduceBegin_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op)
110eb02082bSJunchao Zhang {
111eb02082bSJunchao Zhang   PetscErrorCode         ierr;
112cd620004SJunchao Zhang   PetscSFLink            link;
113dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv     *dat = (PetscSF_Allgatherv*)sf->data;
114dd5b3ca6SJunchao Zhang   PetscInt               rstart;
115cd620004SJunchao Zhang   PetscMPIInt            rank,count,recvcount;
116dd5b3ca6SJunchao Zhang   MPI_Comm               comm;
117cd620004SJunchao Zhang   void                   *rootbuf = NULL,*leafbuf = NULL;
118cd620004SJunchao Zhang   MPI_Request            *req;
119dd5b3ca6SJunchao Zhang 
120dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
121cd620004SJunchao Zhang   ierr = PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,PETSCSF_REDUCE,&link);CHKERRQ(ierr);
12283df288dSJunchao Zhang   if (op == MPI_REPLACE) {
123cd620004SJunchao Zhang     /* REPLACE is only meaningful when all processes have the same leafdata to reduce. Therefore copying from local leafdata is fine */
124dd5b3ca6SJunchao Zhang     ierr = PetscLayoutGetRange(sf->map,&rstart,NULL);CHKERRQ(ierr);
12520c24465SJunchao Zhang     ierr = (*link->Memcpy)(link,rootmtype,rootdata,leafmtype,(const char*)leafdata+(size_t)rstart*link->unitbytes,(size_t)sf->nroots*link->unitbytes);CHKERRQ(ierr);
12671438e86SJunchao Zhang     if (PetscMemTypeDevice(leafmtype) && PetscMemTypeHost(rootmtype)) {ierr = (*link->SyncStream)(link);CHKERRQ(ierr);}
127dd5b3ca6SJunchao Zhang   } else {
128cd620004SJunchao Zhang     /* Reduce leafdata, then scatter to rootdata */
129cd620004SJunchao Zhang     ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr);
130ffc4695bSBarry Smith     ierr = MPI_Comm_rank(comm,&rank);CHKERRMPI(ierr);
131cd620004SJunchao Zhang     ierr = PetscSFLinkPackLeafData(sf,link,PETSCSF_REMOTE,leafdata);CHKERRQ(ierr);
13271438e86SJunchao Zhang     ierr = PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf,link,PETSC_TRUE/* device2host before sending */);CHKERRQ(ierr);
133cd620004SJunchao Zhang     ierr = PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_LEAF2ROOT,&rootbuf,&leafbuf,&req,NULL);CHKERRQ(ierr);
134cd620004SJunchao Zhang     ierr = PetscMPIIntCast(dat->rootbuflen[PETSCSF_REMOTE],&recvcount);CHKERRQ(ierr);
135cd620004SJunchao Zhang     /* Allocate a separate leaf buffer on rank 0 */
136dd400576SPatrick Sanan     if (rank == 0 && !link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi]) {
13720c24465SJunchao Zhang       ierr = PetscSFMalloc(sf,link->leafmtype_mpi,sf->leafbuflen[PETSCSF_REMOTE]*link->unitbytes,(void**)&link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi]);CHKERRQ(ierr);
138dd5b3ca6SJunchao Zhang     }
139cd620004SJunchao Zhang     /* In case we already copied leafdata from device to host (i.e., no use_gpu_aware_mpi), we need to adjust leafbuf on rank 0 */
140dd400576SPatrick Sanan     if (rank == 0 && link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi] == leafbuf) leafbuf = MPI_IN_PLACE;
141cd620004SJunchao Zhang     ierr = PetscMPIIntCast(sf->nleaves*link->bs,&count);CHKERRQ(ierr);
14271438e86SJunchao Zhang     ierr = PetscSFLinkSyncStreamBeforeCallMPI(sf,link,PETSCSF_LEAF2ROOT);CHKERRQ(ierr);
14371438e86SJunchao Zhang     ierr = MPI_Reduce(leafbuf,link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi],count,link->basicunit,op,0,comm);CHKERRMPI(ierr); /* Must do reduce with MPI builltin datatype basicunit */
14471438e86SJunchao Zhang     ierr = MPIU_Iscatterv(link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi],dat->recvcounts,dat->displs,unit,rootbuf,recvcount,unit,0,comm,req);CHKERRMPI(ierr);
145dd5b3ca6SJunchao Zhang   }
146eb02082bSJunchao Zhang   PetscFunctionReturn(0);
147eb02082bSJunchao Zhang }
148eb02082bSJunchao Zhang 
1499319200aSJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFReduceEnd_Allgatherv(PetscSF sf,MPI_Datatype unit,const void *leafdata,void *rootdata,MPI_Op op)
1509319200aSJunchao Zhang {
1519319200aSJunchao Zhang   PetscErrorCode        ierr;
1529319200aSJunchao Zhang   PetscSFLink           link;
1539319200aSJunchao Zhang 
1549319200aSJunchao Zhang   PetscFunctionBegin;
1559319200aSJunchao Zhang   if (op == MPI_REPLACE) {
1569319200aSJunchao Zhang     /* A rare case happens when op is MPI_REPLACE, using GPUs but no GPU aware MPI. In PetscSFReduceBegin_Allgather(v),
1579319200aSJunchao Zhang       we did a device to device copy and in effect finished the communication. But in PetscSFLinkFinishCommunication()
1589319200aSJunchao Zhang       of PetscSFReduceEnd_Basic(), it thinks since there is rootbuf, it calls PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI().
1599319200aSJunchao Zhang       It does a host to device memory copy on rootbuf, wrongly overwritting the results. So we don't overload
1609319200aSJunchao Zhang       PetscSFReduceEnd_Basic() in this case, and just reclaim the link.
1619319200aSJunchao Zhang      */
1629319200aSJunchao Zhang     ierr = PetscSFLinkGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
1639319200aSJunchao Zhang     ierr = PetscSFLinkReclaim(sf,&link);CHKERRQ(ierr);
1649319200aSJunchao Zhang   } else {
1659319200aSJunchao Zhang     ierr = PetscSFReduceEnd_Basic(sf,unit,leafdata,rootdata,op);CHKERRQ(ierr);
1669319200aSJunchao Zhang   }
1679319200aSJunchao Zhang   PetscFunctionReturn(0);
1689319200aSJunchao Zhang }
1699319200aSJunchao Zhang 
170eb02082bSJunchao Zhang static PetscErrorCode PetscSFBcastToZero_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata)
171eb02082bSJunchao Zhang {
172eb02082bSJunchao Zhang   PetscErrorCode         ierr;
173cd620004SJunchao Zhang   PetscSFLink            link;
174855db38dSJunchao Zhang   PetscMPIInt            rank;
175eb02082bSJunchao Zhang 
176eb02082bSJunchao Zhang   PetscFunctionBegin;
177ad227feaSJunchao Zhang   ierr = PetscSFBcastBegin_Gatherv(sf,unit,rootmtype,rootdata,leafmtype,leafdata,MPI_REPLACE);CHKERRQ(ierr);
178cd620004SJunchao Zhang   ierr = PetscSFLinkGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
17971438e86SJunchao Zhang   ierr = PetscSFLinkFinishCommunication(sf,link,PETSCSF_ROOT2LEAF);CHKERRQ(ierr);
180ffc4695bSBarry Smith   ierr = MPI_Comm_rank(PetscObjectComm((PetscObject)sf),&rank);CHKERRMPI(ierr);
181dd400576SPatrick Sanan   if (rank == 0 && PetscMemTypeDevice(leafmtype) && !sf->use_gpu_aware_mpi) {
18220c24465SJunchao Zhang     ierr = (*link->Memcpy)(link,PETSC_MEMTYPE_DEVICE,leafdata,PETSC_MEMTYPE_HOST,link->leafbuf[PETSC_MEMTYPE_HOST],sf->leafbuflen[PETSCSF_REMOTE]*link->unitbytes);CHKERRQ(ierr);
183855db38dSJunchao Zhang   }
184cd620004SJunchao Zhang   ierr = PetscSFLinkReclaim(sf,&link);CHKERRQ(ierr);
185dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
186dd5b3ca6SJunchao Zhang }
187dd5b3ca6SJunchao Zhang 
188dd5b3ca6SJunchao Zhang /* This routine is very tricky (I believe it is rarely used with this kind of graph so just provide a simple but not-optimal implementation).
189dd5b3ca6SJunchao Zhang 
190dd5b3ca6SJunchao Zhang    Suppose we have three ranks. Rank 0 has a root with value 1. Rank 0,1,2 has a leaf with value 2,3,4 respectively. The leaves are connected
191dd5b3ca6SJunchao Zhang    to the root on rank 0. Suppose op=MPI_SUM and rank 0,1,2 gets root state in their rank order. By definition of this routine, rank 0 sees 1
192dd5b3ca6SJunchao Zhang    in root, fetches it into its leafupate, then updates root to 1 + 2 = 3; rank 1 sees 3 in root, fetches it into its leafupate, then updates
193dd5b3ca6SJunchao Zhang    root to 3 + 3 = 6; rank 2 sees 6 in root, fetches it into its leafupdate, then updates root to 6 + 4 = 10.  At the end, leafupdate on rank
194dd5b3ca6SJunchao Zhang    0,1,2 is 1,3,6 respectively. root is 10.
195dd5b3ca6SJunchao Zhang 
196dd5b3ca6SJunchao Zhang    We use a simpler implementation. From the same initial state, we copy leafdata to leafupdate
197dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
198dd5b3ca6SJunchao Zhang         Root     1
199dd5b3ca6SJunchao Zhang         Leaf     2       3         4
200dd5b3ca6SJunchao Zhang      Leafupdate  2       3         4
201dd5b3ca6SJunchao Zhang 
202dd5b3ca6SJunchao Zhang    Do MPI_Exscan on leafupdate,
203dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
204dd5b3ca6SJunchao Zhang         Root     1
205dd5b3ca6SJunchao Zhang         Leaf     2       3         4
206dd5b3ca6SJunchao Zhang      Leafupdate  2       2         5
207dd5b3ca6SJunchao Zhang 
208dd5b3ca6SJunchao Zhang    BcastAndOp from root to leafupdate,
209dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
210dd5b3ca6SJunchao Zhang         Root     1
211dd5b3ca6SJunchao Zhang         Leaf     2       3         4
212dd5b3ca6SJunchao Zhang      Leafupdate  3       3         6
213dd5b3ca6SJunchao Zhang 
214dd5b3ca6SJunchao Zhang    Copy root to leafupdate on rank-0
215dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
216dd5b3ca6SJunchao Zhang         Root     1
217dd5b3ca6SJunchao Zhang         Leaf     2       3         4
218dd5b3ca6SJunchao Zhang      Leafupdate  1       3         6
219dd5b3ca6SJunchao Zhang 
220dd5b3ca6SJunchao Zhang    Reduce from leaf to root,
221dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
222dd5b3ca6SJunchao Zhang         Root     10
223dd5b3ca6SJunchao Zhang         Leaf     2       3         4
224dd5b3ca6SJunchao Zhang      Leafupdate  1       3         6
225dd5b3ca6SJunchao Zhang */
226eb02082bSJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFFetchAndOpBegin_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,void *rootdata,PetscMemType leafmtype,const void *leafdata,void *leafupdate,MPI_Op op)
227dd5b3ca6SJunchao Zhang {
228dd5b3ca6SJunchao Zhang   PetscErrorCode         ierr;
229cd620004SJunchao Zhang   PetscSFLink            link;
230dd5b3ca6SJunchao Zhang   MPI_Comm               comm;
231dd5b3ca6SJunchao Zhang   PetscMPIInt            count;
232dd5b3ca6SJunchao Zhang 
233dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
234855db38dSJunchao Zhang   ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr);
235*2c71b3e2SJacob Faibussowitsch   PetscCheckFalse(PetscMemTypeDevice(rootmtype) || PetscMemTypeDevice(leafmtype),comm,PETSC_ERR_SUP,"Do FetchAndOp on device");
236dd5b3ca6SJunchao Zhang   /* Copy leafdata to leafupdate */
237cd620004SJunchao Zhang   ierr = PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,PETSCSF_FETCH,&link);CHKERRQ(ierr);
238cd620004SJunchao Zhang   ierr = PetscSFLinkPackLeafData(sf,link,PETSCSF_REMOTE,leafdata);CHKERRQ(ierr); /* Sync the device */
23920c24465SJunchao Zhang   ierr = (*link->Memcpy)(link,leafmtype,leafupdate,leafmtype,leafdata,sf->nleaves*link->unitbytes);CHKERRQ(ierr);
240cd620004SJunchao Zhang   ierr = PetscSFLinkGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
241dd5b3ca6SJunchao Zhang 
242dd5b3ca6SJunchao Zhang   /* Exscan on leafupdate and then BcastAndOp rootdata to leafupdate */
24383df288dSJunchao Zhang   if (op == MPI_REPLACE) {
244dd5b3ca6SJunchao Zhang     PetscMPIInt size,rank,prev,next;
245ffc4695bSBarry Smith     ierr = MPI_Comm_rank(comm,&rank);CHKERRMPI(ierr);
246ffc4695bSBarry Smith     ierr = MPI_Comm_size(comm,&size);CHKERRMPI(ierr);
247dd5b3ca6SJunchao Zhang     prev = rank ?            rank-1 : MPI_PROC_NULL;
248dd5b3ca6SJunchao Zhang     next = (rank < size-1) ? rank+1 : MPI_PROC_NULL;
249cd620004SJunchao Zhang     ierr = PetscMPIIntCast(sf->nleaves,&count);CHKERRQ(ierr);
250ffc4695bSBarry Smith     ierr = MPI_Sendrecv_replace(leafupdate,count,unit,next,link->tag,prev,link->tag,comm,MPI_STATUSES_IGNORE);CHKERRMPI(ierr);
251cd620004SJunchao Zhang   } else {
252cd620004SJunchao Zhang     ierr = PetscMPIIntCast(sf->nleaves*link->bs,&count);CHKERRQ(ierr);
253ffc4695bSBarry Smith     ierr = MPI_Exscan(MPI_IN_PLACE,leafupdate,count,link->basicunit,op,comm);CHKERRMPI(ierr);
254cd620004SJunchao Zhang   }
255cd620004SJunchao Zhang   ierr = PetscSFLinkReclaim(sf,&link);CHKERRQ(ierr);
256ad227feaSJunchao Zhang   ierr = PetscSFBcastBegin(sf,unit,rootdata,leafupdate,op);CHKERRQ(ierr);
257ad227feaSJunchao Zhang   ierr = PetscSFBcastEnd(sf,unit,rootdata,leafupdate,op);CHKERRQ(ierr);
258dd5b3ca6SJunchao Zhang 
259dd5b3ca6SJunchao Zhang   /* Bcast roots to rank 0's leafupdate */
260dd5b3ca6SJunchao Zhang   ierr = PetscSFBcastToZero_Private(sf,unit,rootdata,leafupdate);CHKERRQ(ierr); /* Using this line makes Allgather SFs able to inherit this routine */
261dd5b3ca6SJunchao Zhang 
262dd5b3ca6SJunchao Zhang   /* Reduce leafdata to rootdata */
263dd5b3ca6SJunchao Zhang   ierr = PetscSFReduceBegin(sf,unit,leafdata,rootdata,op);CHKERRQ(ierr);
264dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
265dd5b3ca6SJunchao Zhang }
266dd5b3ca6SJunchao Zhang 
26700816365SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFFetchAndOpEnd_Allgatherv(PetscSF sf,MPI_Datatype unit,void *rootdata,const void *leafdata,void *leafupdate,MPI_Op op)
268dd5b3ca6SJunchao Zhang {
269dd5b3ca6SJunchao Zhang   PetscErrorCode         ierr;
270dd5b3ca6SJunchao Zhang 
271dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
272dd5b3ca6SJunchao Zhang   ierr = PetscSFReduceEnd(sf,unit,leafdata,rootdata,op);CHKERRQ(ierr);
273dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
274dd5b3ca6SJunchao Zhang }
275dd5b3ca6SJunchao Zhang 
276dd5b3ca6SJunchao Zhang /* Get root ranks accessing my leaves */
277dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFGetRootRanks_Allgatherv(PetscSF sf,PetscInt *nranks,const PetscMPIInt **ranks,const PetscInt **roffset,const PetscInt **rmine,const PetscInt **rremote)
278dd5b3ca6SJunchao Zhang {
279dd5b3ca6SJunchao Zhang   PetscErrorCode ierr;
280dd5b3ca6SJunchao Zhang   PetscInt       i,j,k,size;
281dd5b3ca6SJunchao Zhang   const PetscInt *range;
282dd5b3ca6SJunchao Zhang 
283dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
284dd5b3ca6SJunchao Zhang   /* Lazily construct these large arrays if users really need them for this type of SF. Very likely, they do not */
285dd5b3ca6SJunchao Zhang   if (sf->nranks && !sf->ranks) { /* On rank!=0, sf->nranks=0. The sf->nranks test makes this routine also works for sfgatherv */
286dd5b3ca6SJunchao Zhang     size = sf->nranks;
287dd5b3ca6SJunchao Zhang     ierr = PetscLayoutGetRanges(sf->map,&range);CHKERRQ(ierr);
288dd5b3ca6SJunchao Zhang     ierr = PetscMalloc4(size,&sf->ranks,size+1,&sf->roffset,sf->nleaves,&sf->rmine,sf->nleaves,&sf->rremote);CHKERRQ(ierr);
289dd5b3ca6SJunchao Zhang     for (i=0; i<size; i++) sf->ranks[i] = i;
290da2e4c71SJunchao Zhang     ierr = PetscArraycpy(sf->roffset,range,size+1);CHKERRQ(ierr);
291dd5b3ca6SJunchao Zhang     for (i=0; i<sf->nleaves; i++) sf->rmine[i] = i; /*rmine are never NULL even for contiguous leaves */
292dd5b3ca6SJunchao Zhang     for (i=0; i<size; i++) {
293dd5b3ca6SJunchao Zhang       for (j=range[i],k=0; j<range[i+1]; j++,k++) sf->rremote[j] = k;
294dd5b3ca6SJunchao Zhang     }
295dd5b3ca6SJunchao Zhang   }
296dd5b3ca6SJunchao Zhang 
297dd5b3ca6SJunchao Zhang   if (nranks)  *nranks  = sf->nranks;
298dd5b3ca6SJunchao Zhang   if (ranks)   *ranks   = sf->ranks;
299dd5b3ca6SJunchao Zhang   if (roffset) *roffset = sf->roffset;
300dd5b3ca6SJunchao Zhang   if (rmine)   *rmine   = sf->rmine;
301dd5b3ca6SJunchao Zhang   if (rremote) *rremote = sf->rremote;
302dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
303dd5b3ca6SJunchao Zhang }
304dd5b3ca6SJunchao Zhang 
305dd5b3ca6SJunchao Zhang /* Get leaf ranks accessing my roots */
306dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFGetLeafRanks_Allgatherv(PetscSF sf,PetscInt *niranks,const PetscMPIInt **iranks,const PetscInt **ioffset,const PetscInt **irootloc)
307dd5b3ca6SJunchao Zhang {
308dd5b3ca6SJunchao Zhang   PetscErrorCode     ierr;
309dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv*)sf->data;
310dd5b3ca6SJunchao Zhang   MPI_Comm           comm;
311dd5b3ca6SJunchao Zhang   PetscMPIInt        size,rank;
312dd5b3ca6SJunchao Zhang   PetscInt           i,j;
313dd5b3ca6SJunchao Zhang 
314dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
315dd5b3ca6SJunchao Zhang   /* Lazily construct these large arrays if users really need them for this type of SF. Very likely, they do not */
316dd5b3ca6SJunchao Zhang   ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr);
317ffc4695bSBarry Smith   ierr = MPI_Comm_size(comm,&size);CHKERRMPI(ierr);
318ffc4695bSBarry Smith   ierr = MPI_Comm_rank(comm,&rank);CHKERRMPI(ierr);
319dd5b3ca6SJunchao Zhang   if (niranks) *niranks = size;
320dd5b3ca6SJunchao Zhang 
321dd5b3ca6SJunchao Zhang   /* PetscSF_Basic has distinguished incoming ranks. Here we do not need that. But we must put self as the first and
322dd5b3ca6SJunchao Zhang      sort other ranks. See comments in PetscSFSetUp_Basic about MatGetBrowsOfAoCols_MPIAIJ on why.
323dd5b3ca6SJunchao Zhang    */
324dd5b3ca6SJunchao Zhang   if (iranks) {
325dd5b3ca6SJunchao Zhang     if (!dat->iranks) {
326dd5b3ca6SJunchao Zhang       ierr = PetscMalloc1(size,&dat->iranks);CHKERRQ(ierr);
327dd5b3ca6SJunchao Zhang       dat->iranks[0] = rank;
328dd5b3ca6SJunchao Zhang       for (i=0,j=1; i<size; i++) {if (i == rank) continue; dat->iranks[j++] = i;}
329dd5b3ca6SJunchao Zhang     }
330dd5b3ca6SJunchao Zhang     *iranks = dat->iranks; /* dat->iranks was init'ed to NULL by PetscNewLog */
331dd5b3ca6SJunchao Zhang   }
332dd5b3ca6SJunchao Zhang 
333dd5b3ca6SJunchao Zhang   if (ioffset) {
334dd5b3ca6SJunchao Zhang     if (!dat->ioffset) {
335dd5b3ca6SJunchao Zhang       ierr = PetscMalloc1(size+1,&dat->ioffset);CHKERRQ(ierr);
336dd5b3ca6SJunchao Zhang       for (i=0; i<=size; i++) dat->ioffset[i] = i*sf->nroots;
337dd5b3ca6SJunchao Zhang     }
338dd5b3ca6SJunchao Zhang     *ioffset = dat->ioffset;
339dd5b3ca6SJunchao Zhang   }
340dd5b3ca6SJunchao Zhang 
341dd5b3ca6SJunchao Zhang   if (irootloc) {
342dd5b3ca6SJunchao Zhang     if (!dat->irootloc) {
343dd5b3ca6SJunchao Zhang       ierr = PetscMalloc1(sf->nleaves,&dat->irootloc);CHKERRQ(ierr);
344dd5b3ca6SJunchao Zhang       for (i=0; i<size; i++) {
345dd5b3ca6SJunchao Zhang         for (j=0; j<sf->nroots; j++) dat->irootloc[i*sf->nroots+j] = j;
346dd5b3ca6SJunchao Zhang       }
347dd5b3ca6SJunchao Zhang     }
348dd5b3ca6SJunchao Zhang     *irootloc = dat->irootloc;
349dd5b3ca6SJunchao Zhang   }
350dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
351dd5b3ca6SJunchao Zhang }
352dd5b3ca6SJunchao Zhang 
353dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFCreateLocalSF_Allgatherv(PetscSF sf,PetscSF *out)
354dd5b3ca6SJunchao Zhang {
355dd5b3ca6SJunchao Zhang   PetscInt       i,nroots,nleaves,rstart,*ilocal;
356dd5b3ca6SJunchao Zhang   PetscSFNode    *iremote;
357dd5b3ca6SJunchao Zhang   PetscSF        lsf;
358dd5b3ca6SJunchao Zhang   PetscErrorCode ierr;
359dd5b3ca6SJunchao Zhang 
360dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
361eb02082bSJunchao Zhang   nleaves = sf->nleaves ? sf->nroots : 0; /* sf->nleaves can be zero with SFGather(v) */
362eb02082bSJunchao Zhang   nroots  = nleaves;
363dd5b3ca6SJunchao Zhang   ierr    = PetscMalloc1(nleaves,&ilocal);CHKERRQ(ierr);
364dd5b3ca6SJunchao Zhang   ierr    = PetscMalloc1(nleaves,&iremote);CHKERRQ(ierr);
365dd5b3ca6SJunchao Zhang   ierr    = PetscLayoutGetRange(sf->map,&rstart,NULL);CHKERRQ(ierr);
366dd5b3ca6SJunchao Zhang 
367dd5b3ca6SJunchao Zhang   for (i=0; i<nleaves; i++) {
368dd5b3ca6SJunchao Zhang     ilocal[i]        = rstart + i; /* lsf does not change leave indices */
369dd5b3ca6SJunchao Zhang     iremote[i].rank  = 0;          /* rank in PETSC_COMM_SELF */
370dd5b3ca6SJunchao Zhang     iremote[i].index = i;          /* root index */
371dd5b3ca6SJunchao Zhang   }
372dd5b3ca6SJunchao Zhang 
373dd5b3ca6SJunchao Zhang   ierr = PetscSFCreate(PETSC_COMM_SELF,&lsf);CHKERRQ(ierr);
374dd5b3ca6SJunchao Zhang   ierr = PetscSFSetGraph(lsf,nroots,nleaves,ilocal,PETSC_OWN_POINTER,iremote,PETSC_OWN_POINTER);CHKERRQ(ierr);
375dd5b3ca6SJunchao Zhang   ierr = PetscSFSetUp(lsf);CHKERRQ(ierr);
376dd5b3ca6SJunchao Zhang   *out = lsf;
377dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
378dd5b3ca6SJunchao Zhang }
379dd5b3ca6SJunchao Zhang 
380dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFCreate_Allgatherv(PetscSF sf)
381dd5b3ca6SJunchao Zhang {
382dd5b3ca6SJunchao Zhang   PetscErrorCode     ierr;
383dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv*)sf->data;
384dd5b3ca6SJunchao Zhang 
385dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
386ad227feaSJunchao Zhang   sf->ops->BcastEnd        = PetscSFBcastEnd_Basic;
3879319200aSJunchao Zhang   sf->ops->ReduceEnd       = PetscSFReduceEnd_Allgatherv;
388cd620004SJunchao Zhang 
389dd5b3ca6SJunchao Zhang   sf->ops->SetUp           = PetscSFSetUp_Allgatherv;
390dd5b3ca6SJunchao Zhang   sf->ops->Reset           = PetscSFReset_Allgatherv;
391dd5b3ca6SJunchao Zhang   sf->ops->Destroy         = PetscSFDestroy_Allgatherv;
392dd5b3ca6SJunchao Zhang   sf->ops->GetRootRanks    = PetscSFGetRootRanks_Allgatherv;
393dd5b3ca6SJunchao Zhang   sf->ops->GetLeafRanks    = PetscSFGetLeafRanks_Allgatherv;
394dd5b3ca6SJunchao Zhang   sf->ops->GetGraph        = PetscSFGetGraph_Allgatherv;
395ad227feaSJunchao Zhang   sf->ops->BcastBegin      = PetscSFBcastBegin_Allgatherv;
396dd5b3ca6SJunchao Zhang   sf->ops->ReduceBegin     = PetscSFReduceBegin_Allgatherv;
397dd5b3ca6SJunchao Zhang   sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Allgatherv;
398dd5b3ca6SJunchao Zhang   sf->ops->FetchAndOpEnd   = PetscSFFetchAndOpEnd_Allgatherv;
399dd5b3ca6SJunchao Zhang   sf->ops->CreateLocalSF   = PetscSFCreateLocalSF_Allgatherv;
400dd5b3ca6SJunchao Zhang   sf->ops->BcastToZero     = PetscSFBcastToZero_Allgatherv;
401dd5b3ca6SJunchao Zhang 
402dd5b3ca6SJunchao Zhang   ierr = PetscNewLog(sf,&dat);CHKERRQ(ierr);
403dd5b3ca6SJunchao Zhang   sf->data = (void*)dat;
404dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
405dd5b3ca6SJunchao Zhang }
406