xref: /petsc/src/vec/is/sf/impls/basic/allgatherv/sfallgatherv.c (revision eb02082b21b627ae704d1b64176a7a4b8db6ed4b)
1dd5b3ca6SJunchao Zhang #include <../src/vec/is/sf/impls/basic/allgatherv/sfallgatherv.h>
2dd5b3ca6SJunchao Zhang 
3*eb02082bSJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFBcastAndOpBegin_Gatherv(PetscSF,MPI_Datatype,PetscMemType,const void*,PetscMemType,void*,MPI_Op);
4dd5b3ca6SJunchao Zhang 
5dd5b3ca6SJunchao Zhang /*===================================================================================*/
6*eb02082bSJunchao Zhang /*              Internal routines for PetscSFPack                                    */
7dd5b3ca6SJunchao Zhang /*===================================================================================*/
8*eb02082bSJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFPackGet_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,const void *leafdata,PetscSFPack *mylink)
9dd5b3ca6SJunchao Zhang {
10dd5b3ca6SJunchao Zhang   PetscErrorCode         ierr;
11*eb02082bSJunchao Zhang   PetscSF_Allgatherv     *dat = (PetscSF_Allgatherv*)sf->data;
12*eb02082bSJunchao Zhang   PetscSFPack            link,*p;
13*eb02082bSJunchao Zhang   PetscBool              match;
14*eb02082bSJunchao Zhang   PetscInt               i,j;
15dd5b3ca6SJunchao Zhang 
16dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
17*eb02082bSJunchao Zhang   ierr = PetscSFPackSetErrorOnUnsupportedOverlap(sf,unit,rootdata,leafdata);CHKERRQ(ierr);
18dd5b3ca6SJunchao Zhang   /* Look for types in cache */
19*eb02082bSJunchao Zhang   for (p=&dat->avail; (link=*p); p=&link->next) {
20dd5b3ca6SJunchao Zhang     ierr = MPIPetsc_Type_compare(unit,link->unit,&match);CHKERRQ(ierr);
21dd5b3ca6SJunchao Zhang     if (match) {
22dd5b3ca6SJunchao Zhang       *p = link->next; /* Remove from available list */
23dd5b3ca6SJunchao Zhang       goto found;
24dd5b3ca6SJunchao Zhang     }
25dd5b3ca6SJunchao Zhang   }
26dd5b3ca6SJunchao Zhang 
27dd5b3ca6SJunchao Zhang   ierr = PetscNew(&link);CHKERRQ(ierr);
28*eb02082bSJunchao Zhang   ierr = PetscSFPackSetUp_Host(sf,link,unit);CHKERRQ(ierr);
29*eb02082bSJunchao Zhang   link->nrootreqs = 1;
30*eb02082bSJunchao Zhang   link->nleafreqs = 0;
31*eb02082bSJunchao Zhang   ierr = PetscMalloc1(4,&link->reqs);CHKERRQ(ierr); /* 4 = (nrootreqs+nleafreqs)*4 */
32*eb02082bSJunchao Zhang   for (i=0; i<4; i++) link->reqs[i] = MPI_REQUEST_NULL; /* Initialized to NULL so that we know which need to be freed in Destroy */
33dd5b3ca6SJunchao Zhang 
34*eb02082bSJunchao Zhang   for (i=0; i<2; i++) {
35*eb02082bSJunchao Zhang     for (j=0; j<2; j++) {
36*eb02082bSJunchao Zhang       link->rootreqs[i][j] = link->reqs + (2*i+j);
37*eb02082bSJunchao Zhang       link->leafreqs[i][j] = NULL; /* leaf requests are not needed. Make it NULL to segfault accident use */
38*eb02082bSJunchao Zhang     }
39*eb02082bSJunchao Zhang   }
40*eb02082bSJunchao Zhang 
41*eb02082bSJunchao Zhang #if defined(PETSC_HAVE_CUDA)
42*eb02082bSJunchao Zhang   link->stream = NULL; /* CUDA default stream */
43*eb02082bSJunchao Zhang #endif
44*eb02082bSJunchao Zhang 
45*eb02082bSJunchao Zhang   /* DO NOT allocate link->rootbuf[]/leafleaf[]. We use lazy allocation since these buffers are likely not needed */
46dd5b3ca6SJunchao Zhang found:
47*eb02082bSJunchao Zhang   link->rootmtype = rootmtype;
48*eb02082bSJunchao Zhang   link->leafmtype = leafmtype;
49*eb02082bSJunchao Zhang #if defined(PETSC_HAVE_CUDA)
50*eb02082bSJunchao Zhang   if (rootmtype == PETSC_MEMTYPE_DEVICE || leafmtype == PETSC_MEMTYPE_DEVICE) {ierr = PetscSFPackSetUp_Device(sf,link,unit);CHKERRQ(ierr);}
51*eb02082bSJunchao Zhang #endif
52*eb02082bSJunchao Zhang   link->rkey      = rootdata;
53*eb02082bSJunchao Zhang   link->lkey      = leafdata;
54dd5b3ca6SJunchao Zhang   link->next      = dat->inuse;
55*eb02082bSJunchao Zhang   dat->inuse      = link;
56dd5b3ca6SJunchao Zhang 
57dd5b3ca6SJunchao Zhang   *mylink         = link;
58dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
59dd5b3ca6SJunchao Zhang }
60dd5b3ca6SJunchao Zhang 
61dd5b3ca6SJunchao Zhang /*===================================================================================*/
62dd5b3ca6SJunchao Zhang /*              Implementations of SF public APIs                                    */
63dd5b3ca6SJunchao Zhang /*===================================================================================*/
64dd5b3ca6SJunchao Zhang 
65dd5b3ca6SJunchao Zhang /* PetscSFGetGraph is non-collective. An implementation should not have collective calls */
66dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFGetGraph_Allgatherv(PetscSF sf,PetscInt *nroots,PetscInt *nleaves,const PetscInt **ilocal,const PetscSFNode **iremote)
67dd5b3ca6SJunchao Zhang {
68dd5b3ca6SJunchao Zhang   PetscErrorCode ierr;
69dd5b3ca6SJunchao Zhang   PetscInt       i,j,k;
70dd5b3ca6SJunchao Zhang   const PetscInt *range;
71dd5b3ca6SJunchao Zhang   PetscMPIInt    size;
72dd5b3ca6SJunchao Zhang 
73dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
74dd5b3ca6SJunchao Zhang   ierr = MPI_Comm_size(PetscObjectComm((PetscObject)sf),&size);CHKERRQ(ierr);
75dd5b3ca6SJunchao Zhang   if (nroots)  *nroots  = sf->nroots;
76dd5b3ca6SJunchao Zhang   if (nleaves) *nleaves = sf->nleaves;
77dd5b3ca6SJunchao Zhang   if (ilocal)  *ilocal  = NULL; /* Contiguous leaves */
78dd5b3ca6SJunchao Zhang   if (iremote) {
79dd5b3ca6SJunchao Zhang     if (!sf->remote && sf->nleaves) { /* The && sf->nleaves makes sfgatherv able to inherit this routine */
80dd5b3ca6SJunchao Zhang       ierr = PetscLayoutGetRanges(sf->map,&range);CHKERRQ(ierr);
81dd5b3ca6SJunchao Zhang       ierr = PetscMalloc1(sf->nleaves,&sf->remote);CHKERRQ(ierr);
82dd5b3ca6SJunchao Zhang       sf->remote_alloc = sf->remote;
83dd5b3ca6SJunchao Zhang       for (i=0; i<size; i++) {
84dd5b3ca6SJunchao Zhang         for (j=range[i],k=0; j<range[i+1]; j++,k++) {
85dd5b3ca6SJunchao Zhang           sf->remote[j].rank  = i;
86dd5b3ca6SJunchao Zhang           sf->remote[j].index = k;
87dd5b3ca6SJunchao Zhang         }
88dd5b3ca6SJunchao Zhang       }
89dd5b3ca6SJunchao Zhang     }
90dd5b3ca6SJunchao Zhang     *iremote = sf->remote;
91dd5b3ca6SJunchao Zhang   }
92dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
93dd5b3ca6SJunchao Zhang }
94dd5b3ca6SJunchao Zhang 
95dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFSetUp_Allgatherv(PetscSF sf)
96dd5b3ca6SJunchao Zhang {
97dd5b3ca6SJunchao Zhang   PetscErrorCode     ierr;
98dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv*)sf->data;
99dd5b3ca6SJunchao Zhang   PetscMPIInt        size;
100dd5b3ca6SJunchao Zhang   PetscInt           i;
101dd5b3ca6SJunchao Zhang   const PetscInt     *range;
102dd5b3ca6SJunchao Zhang 
103dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
104dd5b3ca6SJunchao Zhang   ierr = MPI_Comm_size(PetscObjectComm((PetscObject)sf),&size);CHKERRQ(ierr);
105dd5b3ca6SJunchao Zhang   if (sf->nleaves) { /* This if (sf->nleaves) test makes sfgatherv able to inherit this routine */
106dd5b3ca6SJunchao Zhang     ierr = PetscMalloc1(size,&dat->recvcounts);CHKERRQ(ierr);
107dd5b3ca6SJunchao Zhang     ierr = PetscMalloc1(size,&dat->displs);CHKERRQ(ierr);
108dd5b3ca6SJunchao Zhang     ierr = PetscLayoutGetRanges(sf->map,&range);CHKERRQ(ierr);
109dd5b3ca6SJunchao Zhang 
110dd5b3ca6SJunchao Zhang     for (i=0; i<size; i++) {
111dd5b3ca6SJunchao Zhang       ierr = PetscMPIIntCast(range[i],&dat->displs[i]);CHKERRQ(ierr);
112dd5b3ca6SJunchao Zhang       ierr = PetscMPIIntCast(range[i+1]-range[i],&dat->recvcounts[i]);CHKERRQ(ierr);
113dd5b3ca6SJunchao Zhang     }
114dd5b3ca6SJunchao Zhang   }
115dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
116dd5b3ca6SJunchao Zhang }
117dd5b3ca6SJunchao Zhang 
118dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFReset_Allgatherv(PetscSF sf)
119dd5b3ca6SJunchao Zhang {
120dd5b3ca6SJunchao Zhang   PetscErrorCode         ierr;
121*eb02082bSJunchao Zhang   PetscSF_Allgatherv     *dat = (PetscSF_Allgatherv*)sf->data;
122dd5b3ca6SJunchao Zhang 
123dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
124dd5b3ca6SJunchao Zhang   ierr = PetscFree(dat->iranks);CHKERRQ(ierr);
125dd5b3ca6SJunchao Zhang   ierr = PetscFree(dat->ioffset);CHKERRQ(ierr);
126dd5b3ca6SJunchao Zhang   ierr = PetscFree(dat->irootloc);CHKERRQ(ierr);
127dd5b3ca6SJunchao Zhang   ierr = PetscFree(dat->recvcounts);CHKERRQ(ierr);
128dd5b3ca6SJunchao Zhang   ierr = PetscFree(dat->displs);CHKERRQ(ierr);
129dd5b3ca6SJunchao Zhang   if (dat->inuse) SETERRQ(PetscObjectComm((PetscObject)sf),PETSC_ERR_ARG_WRONGSTATE,"Outstanding operation has not been completed");
130*eb02082bSJunchao Zhang   ierr = PetscSFPackDestoryAvailable(&dat->avail);CHKERRQ(ierr);
131dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
132dd5b3ca6SJunchao Zhang }
133dd5b3ca6SJunchao Zhang 
134dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFDestroy_Allgatherv(PetscSF sf)
135dd5b3ca6SJunchao Zhang {
136dd5b3ca6SJunchao Zhang   PetscErrorCode ierr;
137dd5b3ca6SJunchao Zhang 
138dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
139dd5b3ca6SJunchao Zhang   ierr = PetscSFReset_Allgatherv(sf);CHKERRQ(ierr);
140dd5b3ca6SJunchao Zhang   ierr = PetscFree(sf->data);CHKERRQ(ierr);
141dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
142dd5b3ca6SJunchao Zhang }
143dd5b3ca6SJunchao Zhang 
144*eb02082bSJunchao Zhang static PetscErrorCode PetscSFBcastAndOpBegin_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op)
145dd5b3ca6SJunchao Zhang {
146dd5b3ca6SJunchao Zhang   PetscErrorCode         ierr;
147*eb02082bSJunchao Zhang   PetscSFPack            link;
148dd5b3ca6SJunchao Zhang   PetscMPIInt            sendcount;
149dd5b3ca6SJunchao Zhang   MPI_Comm               comm;
150*eb02082bSJunchao Zhang   char                   *recvbuf;
151dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv     *dat = (PetscSF_Allgatherv*)sf->data;
152dd5b3ca6SJunchao Zhang 
153dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
154*eb02082bSJunchao Zhang   ierr = PetscSFPackGet_Allgatherv(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);CHKERRQ(ierr);
155dd5b3ca6SJunchao Zhang   ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr);
156dd5b3ca6SJunchao Zhang   ierr = PetscMPIIntCast(sf->nroots,&sendcount);CHKERRQ(ierr);
157dd5b3ca6SJunchao Zhang 
158*eb02082bSJunchao Zhang   if (op == MPIU_REPLACE) {recvbuf = (char*)leafdata;}
159*eb02082bSJunchao Zhang   else {
160*eb02082bSJunchao Zhang     if (!link->leafbuf[leafmtype]) {ierr = PetscMallocWithMemType(leafmtype,sf->nleaves*link->unitbytes,(void**)&link->leafbuf[leafmtype]);CHKERRQ(ierr);}
161*eb02082bSJunchao Zhang     recvbuf = link->leafbuf[leafmtype];
162dd5b3ca6SJunchao Zhang   }
163*eb02082bSJunchao Zhang   ierr = MPIU_Iallgatherv(rootdata,sendcount,unit,recvbuf,dat->recvcounts,dat->displs,unit,comm,link->rootreqs[PETSCSF_ROOT2LEAF_BCAST][rootmtype]);CHKERRQ(ierr);
164dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
165dd5b3ca6SJunchao Zhang }
166dd5b3ca6SJunchao Zhang 
167*eb02082bSJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFBcastAndOpEnd_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op)
168dd5b3ca6SJunchao Zhang {
169dd5b3ca6SJunchao Zhang   PetscErrorCode         ierr;
170*eb02082bSJunchao Zhang   PetscSFPack            link;
171*eb02082bSJunchao Zhang 
172*eb02082bSJunchao Zhang   PetscFunctionBegin;
173*eb02082bSJunchao Zhang   ierr = PetscSFPackGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
174*eb02082bSJunchao Zhang   ierr = MPI_Wait(link->rootreqs[PETSCSF_ROOT2LEAF_BCAST][rootmtype],MPI_STATUS_IGNORE);CHKERRQ(ierr);
175*eb02082bSJunchao Zhang   if (op != MPIU_REPLACE) {ierr = PetscSFUnpackAndOpLeafData(sf,link,NULL,leafdata,op,PETSC_FALSE);CHKERRQ(ierr);}
176*eb02082bSJunchao Zhang   ierr = PetscSFPackReclaim(sf,&link);CHKERRQ(ierr);
177*eb02082bSJunchao Zhang   PetscFunctionReturn(0);
178*eb02082bSJunchao Zhang }
179*eb02082bSJunchao Zhang 
180*eb02082bSJunchao Zhang static PetscErrorCode PetscSFReduceBegin_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op)
181*eb02082bSJunchao Zhang {
182*eb02082bSJunchao Zhang   PetscErrorCode         ierr;
183*eb02082bSJunchao Zhang   PetscSFPack            link;
184dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv     *dat = (PetscSF_Allgatherv*)sf->data;
185dd5b3ca6SJunchao Zhang   PetscInt               rstart;
186dd5b3ca6SJunchao Zhang   PetscMPIInt            rank,count,recvcount;
187dd5b3ca6SJunchao Zhang   MPI_Comm               comm;
188dd5b3ca6SJunchao Zhang 
189dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
190*eb02082bSJunchao Zhang   ierr = PetscSFPackGet_Allgatherv(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);CHKERRQ(ierr);
191dd5b3ca6SJunchao Zhang   ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr);
192dd5b3ca6SJunchao Zhang   ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
193dd5b3ca6SJunchao Zhang 
194dd5b3ca6SJunchao Zhang   if (op == MPIU_REPLACE) {
195dd5b3ca6SJunchao Zhang     /* REPLACE is only meaningful when all processes have the same leafdata to readue. Therefore copy from local leafdata is fine */
196dd5b3ca6SJunchao Zhang     ierr = PetscLayoutGetRange(sf->map,&rstart,NULL);CHKERRQ(ierr);
197*eb02082bSJunchao Zhang     ierr = PetscMemcpyWithMemType(rootmtype,leafmtype,rootdata,(const char*)leafdata+(size_t)rstart*link->unitbytes,(size_t)sf->nroots*link->unitbytes);CHKERRQ(ierr);
198dd5b3ca6SJunchao Zhang   } else {
199dd5b3ca6SJunchao Zhang     /* Reduce all leafdata on rank 0, then scatter the result to root buffer, then reduce root buffer to leafdata */
200*eb02082bSJunchao Zhang     if (!rank && !link->leafbuf[leafmtype]) {ierr = PetscMallocWithMemType(leafmtype,sf->nleaves*link->unitbytes,(void**)&link->leafbuf[leafmtype]);CHKERRQ(ierr);}
201dd5b3ca6SJunchao Zhang     ierr = PetscMPIIntCast(sf->nleaves*link->bs,&count);CHKERRQ(ierr);
202dd5b3ca6SJunchao Zhang     ierr = PetscMPIIntCast(sf->nroots,&recvcount);CHKERRQ(ierr);
203*eb02082bSJunchao Zhang     ierr = MPI_Reduce(leafdata,link->leafbuf[leafmtype],count,link->basicunit,op,0,comm);CHKERRQ(ierr); /* Must do reduce with MPI builltin datatype basicunit */
204*eb02082bSJunchao Zhang     if (!link->rootbuf[rootmtype]) {ierr = PetscMallocWithMemType(rootmtype,sf->nroots*link->unitbytes,(void**)&link->rootbuf[rootmtype]);CHKERRQ(ierr);}
205*eb02082bSJunchao Zhang     ierr = MPIU_Iscatterv(link->leafbuf[leafmtype],dat->recvcounts,dat->displs,unit,link->rootbuf[rootmtype],recvcount,unit,0,comm,link->rootreqs[PETSCSF_LEAF2ROOT_REDUCE][rootmtype]);CHKERRQ(ierr);
206dd5b3ca6SJunchao Zhang   }
207dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
208dd5b3ca6SJunchao Zhang }
209dd5b3ca6SJunchao Zhang 
210*eb02082bSJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFReduceEnd_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op)
211dd5b3ca6SJunchao Zhang {
212dd5b3ca6SJunchao Zhang   PetscErrorCode         ierr;
213*eb02082bSJunchao Zhang   PetscSFPack            link;
214dd5b3ca6SJunchao Zhang 
215dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
216*eb02082bSJunchao Zhang   ierr = PetscSFPackGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
217*eb02082bSJunchao Zhang   ierr = MPI_Wait(link->rootreqs[PETSCSF_LEAF2ROOT_REDUCE][rootmtype],MPI_STATUS_IGNORE);CHKERRQ(ierr);
218*eb02082bSJunchao Zhang   if (op != MPIU_REPLACE) {ierr = PetscSFUnpackAndOpRootData(sf,link,NULL,rootdata,op,PETSC_FALSE);CHKERRQ(ierr);}
219*eb02082bSJunchao Zhang   ierr = PetscSFPackReclaim(sf,&link);CHKERRQ(ierr);
220*eb02082bSJunchao Zhang   PetscFunctionReturn(0);
221*eb02082bSJunchao Zhang }
222*eb02082bSJunchao Zhang 
223*eb02082bSJunchao Zhang static PetscErrorCode PetscSFBcastToZero_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata)
224*eb02082bSJunchao Zhang {
225*eb02082bSJunchao Zhang   PetscErrorCode         ierr;
226*eb02082bSJunchao Zhang   PetscSFPack            link;
227*eb02082bSJunchao Zhang 
228*eb02082bSJunchao Zhang   PetscFunctionBegin;
229*eb02082bSJunchao Zhang   ierr = PetscSFBcastAndOpBegin_Gatherv(sf,unit,rootmtype,rootdata,leafmtype,leafdata,MPIU_REPLACE);CHKERRQ(ierr);
230dd5b3ca6SJunchao Zhang   /* A simplified PetscSFBcastAndOpEnd_Allgatherv */
231*eb02082bSJunchao Zhang   ierr = PetscSFPackGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
232*eb02082bSJunchao Zhang   ierr = MPI_Wait(link->rootreqs[PETSCSF_ROOT2LEAF_BCAST][rootmtype],MPI_STATUS_IGNORE);CHKERRQ(ierr);
233*eb02082bSJunchao Zhang   ierr = PetscSFPackReclaim(sf,&link);CHKERRQ(ierr);
234dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
235dd5b3ca6SJunchao Zhang }
236dd5b3ca6SJunchao Zhang 
237dd5b3ca6SJunchao 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).
238dd5b3ca6SJunchao Zhang 
239dd5b3ca6SJunchao 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
240dd5b3ca6SJunchao 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
241dd5b3ca6SJunchao 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
242dd5b3ca6SJunchao 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
243dd5b3ca6SJunchao Zhang    0,1,2 is 1,3,6 respectively. root is 10.
244dd5b3ca6SJunchao Zhang 
245dd5b3ca6SJunchao Zhang    One optimized implementation could be: starting from the initial state:
246dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
247dd5b3ca6SJunchao Zhang         Root     1
248dd5b3ca6SJunchao Zhang         Leaf     2       3         4
249dd5b3ca6SJunchao Zhang 
250dd5b3ca6SJunchao Zhang    Shift leaves rightwards to leafupdate. Rank 0 gathers the root value and puts it in leafupdate. We have:
251dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
252dd5b3ca6SJunchao Zhang         Root     1
253dd5b3ca6SJunchao Zhang         Leaf     2       3         4
254dd5b3ca6SJunchao Zhang      Leafupdate  1       2         3
255dd5b3ca6SJunchao Zhang 
256dd5b3ca6SJunchao Zhang    Then, do MPI_Scan on leafupdate and get:
257dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
258dd5b3ca6SJunchao Zhang         Root     1
259dd5b3ca6SJunchao Zhang         Leaf     2       3         4
260dd5b3ca6SJunchao Zhang      Leafupdate  1       3         6
261dd5b3ca6SJunchao Zhang 
262dd5b3ca6SJunchao Zhang    Rank 2 sums its leaf and leafupdate, scatters the result to the root, and gets
263dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
264dd5b3ca6SJunchao Zhang         Root     10
265dd5b3ca6SJunchao Zhang         Leaf     2       3         4
266dd5b3ca6SJunchao Zhang      Leafupdate  1       3         6
267dd5b3ca6SJunchao Zhang 
268dd5b3ca6SJunchao Zhang    We use a simpler implementation. From the same initial state, we copy leafdata to leafupdate
269dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
270dd5b3ca6SJunchao Zhang         Root     1
271dd5b3ca6SJunchao Zhang         Leaf     2       3         4
272dd5b3ca6SJunchao Zhang      Leafupdate  2       3         4
273dd5b3ca6SJunchao Zhang 
274dd5b3ca6SJunchao Zhang    Do MPI_Exscan on leafupdate,
275dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
276dd5b3ca6SJunchao Zhang         Root     1
277dd5b3ca6SJunchao Zhang         Leaf     2       3         4
278dd5b3ca6SJunchao Zhang      Leafupdate  2       2         5
279dd5b3ca6SJunchao Zhang 
280dd5b3ca6SJunchao Zhang    BcastAndOp from root to leafupdate,
281dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
282dd5b3ca6SJunchao Zhang         Root     1
283dd5b3ca6SJunchao Zhang         Leaf     2       3         4
284dd5b3ca6SJunchao Zhang      Leafupdate  3       3         6
285dd5b3ca6SJunchao Zhang 
286dd5b3ca6SJunchao Zhang    Copy root to leafupdate on rank-0
287dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
288dd5b3ca6SJunchao Zhang         Root     1
289dd5b3ca6SJunchao Zhang         Leaf     2       3         4
290dd5b3ca6SJunchao Zhang      Leafupdate  1       3         6
291dd5b3ca6SJunchao Zhang 
292dd5b3ca6SJunchao Zhang    Reduce from leaf to root,
293dd5b3ca6SJunchao Zhang              rank-0   rank-1    rank-2
294dd5b3ca6SJunchao Zhang         Root     10
295dd5b3ca6SJunchao Zhang         Leaf     2       3         4
296dd5b3ca6SJunchao Zhang      Leafupdate  1       3         6
297dd5b3ca6SJunchao Zhang */
298*eb02082bSJunchao 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)
299dd5b3ca6SJunchao Zhang {
300dd5b3ca6SJunchao Zhang   PetscErrorCode         ierr;
301*eb02082bSJunchao Zhang   PetscSFPack            link;
302dd5b3ca6SJunchao Zhang   MPI_Comm               comm;
303dd5b3ca6SJunchao Zhang   PetscMPIInt            count;
304dd5b3ca6SJunchao Zhang 
305dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
306dd5b3ca6SJunchao Zhang   /* Copy leafdata to leafupdate */
307*eb02082bSJunchao Zhang   ierr = PetscSFPackGet_Allgatherv(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);CHKERRQ(ierr);
308*eb02082bSJunchao Zhang   ierr = PetscMemcpyWithMemType(leafmtype,leafmtype,leafupdate,leafdata,sf->nleaves*link->unitbytes);CHKERRQ(ierr);
309*eb02082bSJunchao Zhang   ierr = PetscSFPackGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr);
310dd5b3ca6SJunchao Zhang 
311dd5b3ca6SJunchao Zhang   /* Exscan on leafupdate and then BcastAndOp rootdata to leafupdate */
312dd5b3ca6SJunchao Zhang   ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr);
313dd5b3ca6SJunchao Zhang   ierr = PetscMPIIntCast(sf->nleaves,&count);CHKERRQ(ierr);
314dd5b3ca6SJunchao Zhang   if (op == MPIU_REPLACE) {
315dd5b3ca6SJunchao Zhang     PetscMPIInt size,rank,prev,next;
316dd5b3ca6SJunchao Zhang     ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
317dd5b3ca6SJunchao Zhang     ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
318dd5b3ca6SJunchao Zhang     prev = rank ?            rank-1 : MPI_PROC_NULL;
319dd5b3ca6SJunchao Zhang     next = (rank < size-1) ? rank+1 : MPI_PROC_NULL;
320dd5b3ca6SJunchao Zhang     ierr = MPI_Sendrecv_replace(leafupdate,count,unit,next,link->tag,prev,link->tag,comm,MPI_STATUSES_IGNORE);CHKERRQ(ierr);
321dd5b3ca6SJunchao Zhang   } else {ierr = MPI_Exscan(MPI_IN_PLACE,leafupdate,count,unit,op,comm);CHKERRQ(ierr);}
322*eb02082bSJunchao Zhang   ierr = PetscSFPackReclaim(sf,&link);CHKERRQ(ierr);
323dd5b3ca6SJunchao Zhang   ierr = PetscSFBcastAndOpBegin(sf,unit,rootdata,leafupdate,op);CHKERRQ(ierr);
324dd5b3ca6SJunchao Zhang   ierr = PetscSFBcastAndOpEnd(sf,unit,rootdata,leafupdate,op);CHKERRQ(ierr);
325dd5b3ca6SJunchao Zhang 
326dd5b3ca6SJunchao Zhang   /* Bcast roots to rank 0's leafupdate */
327dd5b3ca6SJunchao Zhang   ierr = PetscSFBcastToZero_Private(sf,unit,rootdata,leafupdate);CHKERRQ(ierr); /* Using this line makes Allgather SFs able to inherit this routine */
328dd5b3ca6SJunchao Zhang 
329dd5b3ca6SJunchao Zhang   /* Reduce leafdata to rootdata */
330dd5b3ca6SJunchao Zhang   ierr = PetscSFReduceBegin(sf,unit,leafdata,rootdata,op);CHKERRQ(ierr);
331dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
332dd5b3ca6SJunchao Zhang }
333dd5b3ca6SJunchao Zhang 
334*eb02082bSJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFFetchAndOpEnd_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,void *rootdata,PetscMemType leafmtype,const void *leafdata,void *leafupdate,MPI_Op op)
335dd5b3ca6SJunchao Zhang {
336dd5b3ca6SJunchao Zhang   PetscErrorCode         ierr;
337dd5b3ca6SJunchao Zhang 
338dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
339dd5b3ca6SJunchao Zhang   ierr = PetscSFReduceEnd(sf,unit,leafdata,rootdata,op);CHKERRQ(ierr);
340dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
341dd5b3ca6SJunchao Zhang }
342dd5b3ca6SJunchao Zhang 
343dd5b3ca6SJunchao Zhang /* Get root ranks accessing my leaves */
344dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFGetRootRanks_Allgatherv(PetscSF sf,PetscInt *nranks,const PetscMPIInt **ranks,const PetscInt **roffset,const PetscInt **rmine,const PetscInt **rremote)
345dd5b3ca6SJunchao Zhang {
346dd5b3ca6SJunchao Zhang   PetscErrorCode ierr;
347dd5b3ca6SJunchao Zhang   PetscInt       i,j,k,size;
348dd5b3ca6SJunchao Zhang   const PetscInt *range;
349dd5b3ca6SJunchao Zhang 
350dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
351dd5b3ca6SJunchao Zhang   /* Lazily construct these large arrays if users really need them for this type of SF. Very likely, they do not */
352dd5b3ca6SJunchao Zhang   if (sf->nranks && !sf->ranks) { /* On rank!=0, sf->nranks=0. The sf->nranks test makes this routine also works for sfgatherv */
353dd5b3ca6SJunchao Zhang     size = sf->nranks;
354dd5b3ca6SJunchao Zhang     ierr = PetscLayoutGetRanges(sf->map,&range);CHKERRQ(ierr);
355dd5b3ca6SJunchao Zhang     ierr = PetscMalloc4(size,&sf->ranks,size+1,&sf->roffset,sf->nleaves,&sf->rmine,sf->nleaves,&sf->rremote);CHKERRQ(ierr);
356dd5b3ca6SJunchao Zhang     for (i=0; i<size; i++) sf->ranks[i] = i;
357da2e4c71SJunchao Zhang     ierr = PetscArraycpy(sf->roffset,range,size+1);CHKERRQ(ierr);
358dd5b3ca6SJunchao Zhang     for (i=0; i<sf->nleaves; i++) sf->rmine[i] = i; /*rmine are never NULL even for contiguous leaves */
359dd5b3ca6SJunchao Zhang     for (i=0; i<size; i++) {
360dd5b3ca6SJunchao Zhang       for (j=range[i],k=0; j<range[i+1]; j++,k++) sf->rremote[j] = k;
361dd5b3ca6SJunchao Zhang     }
362dd5b3ca6SJunchao Zhang   }
363dd5b3ca6SJunchao Zhang 
364dd5b3ca6SJunchao Zhang   if (nranks)  *nranks  = sf->nranks;
365dd5b3ca6SJunchao Zhang   if (ranks)   *ranks   = sf->ranks;
366dd5b3ca6SJunchao Zhang   if (roffset) *roffset = sf->roffset;
367dd5b3ca6SJunchao Zhang   if (rmine)   *rmine   = sf->rmine;
368dd5b3ca6SJunchao Zhang   if (rremote) *rremote = sf->rremote;
369dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
370dd5b3ca6SJunchao Zhang }
371dd5b3ca6SJunchao Zhang 
372dd5b3ca6SJunchao Zhang /* Get leaf ranks accessing my roots */
373dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFGetLeafRanks_Allgatherv(PetscSF sf,PetscInt *niranks,const PetscMPIInt **iranks,const PetscInt **ioffset,const PetscInt **irootloc)
374dd5b3ca6SJunchao Zhang {
375dd5b3ca6SJunchao Zhang   PetscErrorCode     ierr;
376dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv*)sf->data;
377dd5b3ca6SJunchao Zhang   MPI_Comm           comm;
378dd5b3ca6SJunchao Zhang   PetscMPIInt        size,rank;
379dd5b3ca6SJunchao Zhang   PetscInt           i,j;
380dd5b3ca6SJunchao Zhang 
381dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
382dd5b3ca6SJunchao Zhang   /* Lazily construct these large arrays if users really need them for this type of SF. Very likely, they do not */
383dd5b3ca6SJunchao Zhang   ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr);
384dd5b3ca6SJunchao Zhang   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
385dd5b3ca6SJunchao Zhang   ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
386dd5b3ca6SJunchao Zhang   if (niranks) *niranks = size;
387dd5b3ca6SJunchao Zhang 
388dd5b3ca6SJunchao Zhang   /* PetscSF_Basic has distinguished incoming ranks. Here we do not need that. But we must put self as the first and
389dd5b3ca6SJunchao Zhang      sort other ranks. See comments in PetscSFSetUp_Basic about MatGetBrowsOfAoCols_MPIAIJ on why.
390dd5b3ca6SJunchao Zhang    */
391dd5b3ca6SJunchao Zhang   if (iranks) {
392dd5b3ca6SJunchao Zhang     if (!dat->iranks) {
393dd5b3ca6SJunchao Zhang       ierr = PetscMalloc1(size,&dat->iranks);CHKERRQ(ierr);
394dd5b3ca6SJunchao Zhang       dat->iranks[0] = rank;
395dd5b3ca6SJunchao Zhang       for (i=0,j=1; i<size; i++) {if (i == rank) continue; dat->iranks[j++] = i;}
396dd5b3ca6SJunchao Zhang     }
397dd5b3ca6SJunchao Zhang     *iranks = dat->iranks; /* dat->iranks was init'ed to NULL by PetscNewLog */
398dd5b3ca6SJunchao Zhang   }
399dd5b3ca6SJunchao Zhang 
400dd5b3ca6SJunchao Zhang   if (ioffset) {
401dd5b3ca6SJunchao Zhang     if (!dat->ioffset) {
402dd5b3ca6SJunchao Zhang       ierr = PetscMalloc1(size+1,&dat->ioffset);CHKERRQ(ierr);
403dd5b3ca6SJunchao Zhang       for (i=0; i<=size; i++) dat->ioffset[i] = i*sf->nroots;
404dd5b3ca6SJunchao Zhang     }
405dd5b3ca6SJunchao Zhang     *ioffset = dat->ioffset;
406dd5b3ca6SJunchao Zhang   }
407dd5b3ca6SJunchao Zhang 
408dd5b3ca6SJunchao Zhang   if (irootloc) {
409dd5b3ca6SJunchao Zhang     if (!dat->irootloc) {
410dd5b3ca6SJunchao Zhang       ierr = PetscMalloc1(sf->nleaves,&dat->irootloc);CHKERRQ(ierr);
411dd5b3ca6SJunchao Zhang       for (i=0; i<size; i++) {
412dd5b3ca6SJunchao Zhang         for (j=0; j<sf->nroots; j++) dat->irootloc[i*sf->nroots+j] = j;
413dd5b3ca6SJunchao Zhang       }
414dd5b3ca6SJunchao Zhang     }
415dd5b3ca6SJunchao Zhang     *irootloc = dat->irootloc;
416dd5b3ca6SJunchao Zhang   }
417dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
418dd5b3ca6SJunchao Zhang }
419dd5b3ca6SJunchao Zhang 
420dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFCreateLocalSF_Allgatherv(PetscSF sf,PetscSF *out)
421dd5b3ca6SJunchao Zhang {
422dd5b3ca6SJunchao Zhang   PetscInt       i,nroots,nleaves,rstart,*ilocal;
423dd5b3ca6SJunchao Zhang   PetscSFNode    *iremote;
424dd5b3ca6SJunchao Zhang   PetscSF        lsf;
425dd5b3ca6SJunchao Zhang   PetscErrorCode ierr;
426dd5b3ca6SJunchao Zhang 
427dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
428*eb02082bSJunchao Zhang   nleaves = sf->nleaves ? sf->nroots : 0; /* sf->nleaves can be zero with SFGather(v) */
429*eb02082bSJunchao Zhang   nroots  = nleaves;
430dd5b3ca6SJunchao Zhang   ierr    = PetscMalloc1(nleaves,&ilocal);CHKERRQ(ierr);
431dd5b3ca6SJunchao Zhang   ierr    = PetscMalloc1(nleaves,&iremote);CHKERRQ(ierr);
432dd5b3ca6SJunchao Zhang   ierr    = PetscLayoutGetRange(sf->map,&rstart,NULL);CHKERRQ(ierr);
433dd5b3ca6SJunchao Zhang 
434dd5b3ca6SJunchao Zhang   for (i=0; i<nleaves; i++) {
435dd5b3ca6SJunchao Zhang     ilocal[i]        = rstart + i; /* lsf does not change leave indices */
436dd5b3ca6SJunchao Zhang     iremote[i].rank  = 0;          /* rank in PETSC_COMM_SELF */
437dd5b3ca6SJunchao Zhang     iremote[i].index = i;          /* root index */
438dd5b3ca6SJunchao Zhang   }
439dd5b3ca6SJunchao Zhang 
440dd5b3ca6SJunchao Zhang   ierr = PetscSFCreate(PETSC_COMM_SELF,&lsf);CHKERRQ(ierr);
441dd5b3ca6SJunchao Zhang   ierr = PetscSFSetGraph(lsf,nroots,nleaves,ilocal,PETSC_OWN_POINTER,iremote,PETSC_OWN_POINTER);CHKERRQ(ierr);
442dd5b3ca6SJunchao Zhang   ierr = PetscSFSetUp(lsf);CHKERRQ(ierr);
443dd5b3ca6SJunchao Zhang   *out = lsf;
444dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
445dd5b3ca6SJunchao Zhang }
446dd5b3ca6SJunchao Zhang 
447dd5b3ca6SJunchao Zhang PETSC_INTERN PetscErrorCode PetscSFCreate_Allgatherv(PetscSF sf)
448dd5b3ca6SJunchao Zhang {
449dd5b3ca6SJunchao Zhang   PetscErrorCode     ierr;
450dd5b3ca6SJunchao Zhang   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv*)sf->data;
451dd5b3ca6SJunchao Zhang 
452dd5b3ca6SJunchao Zhang   PetscFunctionBegin;
453dd5b3ca6SJunchao Zhang   sf->ops->SetUp           = PetscSFSetUp_Allgatherv;
454dd5b3ca6SJunchao Zhang   sf->ops->Reset           = PetscSFReset_Allgatherv;
455dd5b3ca6SJunchao Zhang   sf->ops->Destroy         = PetscSFDestroy_Allgatherv;
456dd5b3ca6SJunchao Zhang   sf->ops->GetRootRanks    = PetscSFGetRootRanks_Allgatherv;
457dd5b3ca6SJunchao Zhang   sf->ops->GetLeafRanks    = PetscSFGetLeafRanks_Allgatherv;
458dd5b3ca6SJunchao Zhang   sf->ops->GetGraph        = PetscSFGetGraph_Allgatherv;
459dd5b3ca6SJunchao Zhang   sf->ops->BcastAndOpBegin = PetscSFBcastAndOpBegin_Allgatherv;
460dd5b3ca6SJunchao Zhang   sf->ops->BcastAndOpEnd   = PetscSFBcastAndOpEnd_Allgatherv;
461dd5b3ca6SJunchao Zhang   sf->ops->ReduceBegin     = PetscSFReduceBegin_Allgatherv;
462dd5b3ca6SJunchao Zhang   sf->ops->ReduceEnd       = PetscSFReduceEnd_Allgatherv;
463dd5b3ca6SJunchao Zhang   sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Allgatherv;
464dd5b3ca6SJunchao Zhang   sf->ops->FetchAndOpEnd   = PetscSFFetchAndOpEnd_Allgatherv;
465dd5b3ca6SJunchao Zhang   sf->ops->CreateLocalSF   = PetscSFCreateLocalSF_Allgatherv;
466dd5b3ca6SJunchao Zhang   sf->ops->BcastToZero     = PetscSFBcastToZero_Allgatherv;
467dd5b3ca6SJunchao Zhang 
468dd5b3ca6SJunchao Zhang   ierr = PetscNewLog(sf,&dat);CHKERRQ(ierr);
469dd5b3ca6SJunchao Zhang   sf->data = (void*)dat;
470dd5b3ca6SJunchao Zhang   PetscFunctionReturn(0);
471dd5b3ca6SJunchao Zhang }
472