Actual source code: sfallgatherv.c

petsc-master 2020-07-14
Report Typos and Errors
  1:  #include <../src/vec/is/sf/impls/basic/allgatherv/sfallgatherv.h>

  3: PETSC_INTERN PetscErrorCode PetscSFBcastAndOpBegin_Gatherv(PetscSF,MPI_Datatype,PetscMemType,const void*,PetscMemType,void*,MPI_Op);

  5: /* PetscSFGetGraph is non-collective. An implementation should not have collective calls */
  6: PETSC_INTERN PetscErrorCode PetscSFGetGraph_Allgatherv(PetscSF sf,PetscInt *nroots,PetscInt *nleaves,const PetscInt **ilocal,const PetscSFNode **iremote)
  7: {
  9:   PetscInt       i,j,k;
 10:   const PetscInt *range;
 11:   PetscMPIInt    size;

 14:   MPI_Comm_size(PetscObjectComm((PetscObject)sf),&size);
 15:   if (nroots)  *nroots  = sf->nroots;
 16:   if (nleaves) *nleaves = sf->nleaves;
 17:   if (ilocal)  *ilocal  = NULL; /* Contiguous leaves */
 18:   if (iremote) {
 19:     if (!sf->remote && sf->nleaves) { /* The && sf->nleaves makes sfgatherv able to inherit this routine */
 20:       PetscLayoutGetRanges(sf->map,&range);
 21:       PetscMalloc1(sf->nleaves,&sf->remote);
 22:       sf->remote_alloc = sf->remote;
 23:       for (i=0; i<size; i++) {
 24:         for (j=range[i],k=0; j<range[i+1]; j++,k++) {
 25:           sf->remote[j].rank  = i;
 26:           sf->remote[j].index = k;
 27:         }
 28:       }
 29:     }
 30:     *iremote = sf->remote;
 31:   }
 32:   return(0);
 33: }

 35: PETSC_INTERN PetscErrorCode PetscSFSetUp_Allgatherv(PetscSF sf)
 36: {
 37:   PetscErrorCode     ierr;
 38:   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv*)sf->data;
 39:   PetscMPIInt        size;
 40:   PetscInt           i;
 41:   const PetscInt     *range;

 44:   PetscSFSetUp_Allgather(sf);
 45:   MPI_Comm_size(PetscObjectComm((PetscObject)sf),&size);
 46:   if (sf->nleaves) { /* This if (sf->nleaves) test makes sfgatherv able to inherit this routine */
 47:     PetscMalloc1(size,&dat->recvcounts);
 48:     PetscMalloc1(size,&dat->displs);
 49:     PetscLayoutGetRanges(sf->map,&range);

 51:     for (i=0; i<size; i++) {
 52:       PetscMPIIntCast(range[i],&dat->displs[i]);
 53:       PetscMPIIntCast(range[i+1]-range[i],&dat->recvcounts[i]);
 54:     }
 55:   }
 56:   return(0);
 57: }

 59: PETSC_INTERN PetscErrorCode PetscSFReset_Allgatherv(PetscSF sf)
 60: {
 61:   PetscErrorCode         ierr;
 62:   PetscSF_Allgatherv     *dat = (PetscSF_Allgatherv*)sf->data;

 65:   PetscFree(dat->iranks);
 66:   PetscFree(dat->ioffset);
 67:   PetscFree(dat->irootloc);
 68:   PetscFree(dat->recvcounts);
 69:   PetscFree(dat->displs);
 70:   if (dat->inuse) SETERRQ(PetscObjectComm((PetscObject)sf),PETSC_ERR_ARG_WRONGSTATE,"Outstanding operation has not been completed");
 71:   PetscSFLinkDestroy(sf,&dat->avail);
 72:   return(0);
 73: }

 75: PETSC_INTERN PetscErrorCode PetscSFDestroy_Allgatherv(PetscSF sf)
 76: {

 80:   PetscSFReset_Allgatherv(sf);
 81:   PetscFree(sf->data);
 82:   return(0);
 83: }

 85: static PetscErrorCode PetscSFBcastAndOpBegin_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op)
 86: {
 87:   PetscErrorCode         ierr;
 88:   PetscSFLink            link;
 89:   PetscMPIInt            sendcount;
 90:   MPI_Comm               comm;
 91:   void                   *rootbuf = NULL,*leafbuf = NULL;
 92:   MPI_Request            *req;
 93:   PetscSF_Allgatherv     *dat = (PetscSF_Allgatherv*)sf->data;

 96:   PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,PETSCSF_BCAST,&link);
 97:   PetscSFLinkPackRootData(sf,link,PETSCSF_REMOTE,rootdata);
 98:   PetscObjectGetComm((PetscObject)sf,&comm);
 99:   PetscMPIIntCast(sf->nroots,&sendcount);
100:   PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_../../../../../..2LEAF,&rootbuf,&leafbuf,&req,NULL);
101:   MPIU_Iallgatherv(rootbuf,sendcount,unit,leafbuf,dat->recvcounts,dat->displs,unit,comm,req);
102:   return(0);
103: }

105: static PetscErrorCode PetscSFReduceBegin_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op)
106: {
107:   PetscErrorCode         ierr;
108:   PetscSFLink            link;
109:   PetscSF_Allgatherv     *dat = (PetscSF_Allgatherv*)sf->data;
110:   PetscInt               rstart;
111:   PetscMPIInt            rank,count,recvcount;
112:   MPI_Comm               comm;
113:   void                   *rootbuf = NULL,*leafbuf = NULL;
114:   MPI_Request            *req;

117:   PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,PETSCSF_REDUCE,&link);
118:   if (op == MPIU_REPLACE) {
119:     /* REPLACE is only meaningful when all processes have the same leafdata to reduce. Therefore copying from local leafdata is fine */
120:     PetscLayoutGetRange(sf->map,&rstart,NULL);
121:     PetscSFLinkMemcpy(sf,link,rootmtype,rootdata,leafmtype,(const char*)leafdata+(size_t)rstart*link->unitbytes,(size_t)sf->nroots*link->unitbytes);
122:   } else {
123:     /* Reduce leafdata, then scatter to rootdata */
124:     PetscObjectGetComm((PetscObject)sf,&comm);
125:     MPI_Comm_rank(comm,&rank);
126:     PetscSFLinkPackLeafData(sf,link,PETSCSF_REMOTE,leafdata);
127:     PetscSFLinkGetMPIBuffersAndRequests(sf,link,PETSCSF_LEAF2../../../../../..,&rootbuf,&leafbuf,&req,NULL);
128:     PetscMPIIntCast(dat->rootbuflen[PETSCSF_REMOTE],&recvcount);
129:     /* Allocate a separate leaf buffer on rank 0 */
130:     if (!rank && !link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi]) {
131:       PetscMallocWithMemType(link->leafmtype_mpi,sf->leafbuflen[PETSCSF_REMOTE]*link->unitbytes,(void**)&link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi]);
132:     }
133:     /* 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 */
134:     if (!rank && link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi] == leafbuf) leafbuf = MPI_IN_PLACE;
135:     PetscMPIIntCast(sf->nleaves*link->bs,&count);
136:     MPI_Reduce(leafbuf,link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi],count,link->basicunit,op,0,comm); /* Must do reduce with MPI builltin datatype basicunit */
137:     MPIU_Iscatterv(link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi],dat->recvcounts,dat->displs,unit,rootbuf,recvcount,unit,0,comm,req);
138:   }
139:   return(0);
140: }

142: static PetscErrorCode PetscSFBcastToZero_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata)
143: {
144:   PetscErrorCode         ierr;
145:   PetscSFLink            link;
146:   PetscMPIInt            rank;

149:   PetscSFBcastAndOpBegin_Gatherv(sf,unit,rootmtype,rootdata,leafmtype,leafdata,MPIU_REPLACE);
150:   PetscSFLinkGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);
151:   PetscSFLinkMPIWaitall(sf,link,PETSCSF_../../../../../..2LEAF);
152:   MPI_Comm_rank(PetscObjectComm((PetscObject)sf),&rank);
153:   if (!rank && leafmtype == PETSC_MEMTYPE_DEVICE && !sf->use_gpu_aware_mpi) {
154:     PetscSFLinkMemcpy(sf,link,PETSC_MEMTYPE_DEVICE,leafdata,PETSC_MEMTYPE_HOST,link->leafbuf[PETSC_MEMTYPE_HOST],sf->leafbuflen[PETSCSF_REMOTE]*link->unitbytes);
155:   }
156:   PetscSFLinkReclaim(sf,&link);
157:   return(0);
158: }

160: /* 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).

162:    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
163:    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
164:    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
165:    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
166:    0,1,2 is 1,3,6 respectively. root is 10.

168:    We use a simpler implementation. From the same initial state, we copy leafdata to leafupdate
169:              rank-0   rank-1    rank-2
170:         Root     1
171:         Leaf     2       3         4
172:      Leafupdate  2       3         4

174:    Do MPI_Exscan on leafupdate,
175:              rank-0   rank-1    rank-2
176:         Root     1
177:         Leaf     2       3         4
178:      Leafupdate  2       2         5

180:    BcastAndOp from root to leafupdate,
181:              rank-0   rank-1    rank-2
182:         Root     1
183:         Leaf     2       3         4
184:      Leafupdate  3       3         6

186:    Copy root to leafupdate on rank-0
187:              rank-0   rank-1    rank-2
188:         Root     1
189:         Leaf     2       3         4
190:      Leafupdate  1       3         6

192:    Reduce from leaf to root,
193:              rank-0   rank-1    rank-2
194:         Root     10
195:         Leaf     2       3         4
196:      Leafupdate  1       3         6
197: */
198: PETSC_INTERN PetscErrorCode PetscSFFetchAndOpBegin_Allgatherv(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,void *rootdata,PetscMemType leafmtype,const void *leafdata,void *leafupdate,MPI_Op op)
199: {
200:   PetscErrorCode         ierr;
201:   PetscSFLink            link;
202:   MPI_Comm               comm;
203:   PetscMPIInt            count;

206:   PetscObjectGetComm((PetscObject)sf,&comm);
207:   if (!sf->use_gpu_aware_mpi && (rootmtype == PETSC_MEMTYPE_DEVICE || leafmtype == PETSC_MEMTYPE_DEVICE)) SETERRQ(comm,PETSC_ERR_SUP,"Do FetchAndOp on device but not use gpu-aware MPI");
208:   /* Copy leafdata to leafupdate */
209:   PetscSFLinkCreate(sf,unit,rootmtype,rootdata,leafmtype,leafdata,op,PETSCSF_FETCH,&link);
210:   PetscSFLinkPackLeafData(sf,link,PETSCSF_REMOTE,leafdata); /* Sync the device */
211:   PetscSFLinkMemcpy(sf,link,leafmtype,leafupdate,leafmtype,leafdata,sf->nleaves*link->unitbytes);
212:   PetscSFLinkGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);

214:   /* Exscan on leafupdate and then BcastAndOp rootdata to leafupdate */
215:   if (op == MPIU_REPLACE) {
216:     PetscMPIInt size,rank,prev,next;
217:     MPI_Comm_rank(comm,&rank);
218:     MPI_Comm_size(comm,&size);
219:     prev = rank ?            rank-1 : MPI_PROC_NULL;
220:     next = (rank < size-1) ? rank+1 : MPI_PROC_NULL;
221:     PetscMPIIntCast(sf->nleaves,&count);
222:     MPI_Sendrecv_replace(leafupdate,count,unit,next,link->tag,prev,link->tag,comm,MPI_STATUSES_IGNORE);
223:   } else {
224:     PetscMPIIntCast(sf->nleaves*link->bs,&count);
225:     MPI_Exscan(MPI_IN_PLACE,leafupdate,count,link->basicunit,op,comm);
226:   }
227:   PetscSFLinkReclaim(sf,&link);
228:   PetscSFBcastAndOpBegin(sf,unit,rootdata,leafupdate,op);
229:   PetscSFBcastAndOpEnd(sf,unit,rootdata,leafupdate,op);

231:   /* Bcast roots to rank 0's leafupdate */
232:   PetscSFBcastToZero_Private(sf,unit,rootdata,leafupdate); /* Using this line makes Allgather SFs able to inherit this routine */

234:   /* Reduce leafdata to rootdata */
235:   PetscSFReduceBegin(sf,unit,leafdata,rootdata,op);
236:   return(0);
237: }

239: PETSC_INTERN PetscErrorCode PetscSFFetchAndOpEnd_Allgatherv(PetscSF sf,MPI_Datatype unit,void *rootdata,const void *leafdata,void *leafupdate,MPI_Op op)
240: {
241:   PetscErrorCode         ierr;

244:   PetscSFReduceEnd(sf,unit,leafdata,rootdata,op);
245:   return(0);
246: }

248: /* Get root ranks accessing my leaves */
249: PETSC_INTERN PetscErrorCode PetscSFGetRootRanks_Allgatherv(PetscSF sf,PetscInt *nranks,const PetscMPIInt **ranks,const PetscInt **roffset,const PetscInt **rmine,const PetscInt **rremote)
250: {
252:   PetscInt       i,j,k,size;
253:   const PetscInt *range;

256:   /* Lazily construct these large arrays if users really need them for this type of SF. Very likely, they do not */
257:   if (sf->nranks && !sf->ranks) { /* On rank!=0, sf->nranks=0. The sf->nranks test makes this routine also works for sfgatherv */
258:     size = sf->nranks;
259:     PetscLayoutGetRanges(sf->map,&range);
260:     PetscMalloc4(size,&sf->ranks,size+1,&sf->roffset,sf->nleaves,&sf->rmine,sf->nleaves,&sf->rremote);
261:     for (i=0; i<size; i++) sf->ranks[i] = i;
262:     PetscArraycpy(sf->roffset,range,size+1);
263:     for (i=0; i<sf->nleaves; i++) sf->rmine[i] = i; /*rmine are never NULL even for contiguous leaves */
264:     for (i=0; i<size; i++) {
265:       for (j=range[i],k=0; j<range[i+1]; j++,k++) sf->rremote[j] = k;
266:     }
267:   }

269:   if (nranks)  *nranks  = sf->nranks;
270:   if (ranks)   *ranks   = sf->ranks;
271:   if (roffset) *roffset = sf->roffset;
272:   if (rmine)   *rmine   = sf->rmine;
273:   if (rremote) *rremote = sf->rremote;
274:   return(0);
275: }

277: /* Get leaf ranks accessing my roots */
278: PETSC_INTERN PetscErrorCode PetscSFGetLeafRanks_Allgatherv(PetscSF sf,PetscInt *niranks,const PetscMPIInt **iranks,const PetscInt **ioffset,const PetscInt **irootloc)
279: {
280:   PetscErrorCode     ierr;
281:   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv*)sf->data;
282:   MPI_Comm           comm;
283:   PetscMPIInt        size,rank;
284:   PetscInt           i,j;

287:   /* Lazily construct these large arrays if users really need them for this type of SF. Very likely, they do not */
288:   PetscObjectGetComm((PetscObject)sf,&comm);
289:   MPI_Comm_size(comm,&size);
290:   MPI_Comm_rank(comm,&rank);
291:   if (niranks) *niranks = size;

293:   /* PetscSF_Basic has distinguished incoming ranks. Here we do not need that. But we must put self as the first and
294:      sort other ranks. See comments in PetscSFSetUp_Basic about MatGetBrowsOfAoCols_MPIAIJ on why.
295:    */
296:   if (iranks) {
297:     if (!dat->iranks) {
298:       PetscMalloc1(size,&dat->iranks);
299:       dat->iranks[0] = rank;
300:       for (i=0,j=1; i<size; i++) {if (i == rank) continue; dat->iranks[j++] = i;}
301:     }
302:     *iranks = dat->iranks; /* dat->iranks was init'ed to NULL by PetscNewLog */
303:   }

305:   if (ioffset) {
306:     if (!dat->ioffset) {
307:       PetscMalloc1(size+1,&dat->ioffset);
308:       for (i=0; i<=size; i++) dat->ioffset[i] = i*sf->nroots;
309:     }
310:     *ioffset = dat->ioffset;
311:   }

313:   if (irootloc) {
314:     if (!dat->irootloc) {
315:       PetscMalloc1(sf->nleaves,&dat->irootloc);
316:       for (i=0; i<size; i++) {
317:         for (j=0; j<sf->nroots; j++) dat->irootloc[i*sf->nroots+j] = j;
318:       }
319:     }
320:     *irootloc = dat->irootloc;
321:   }
322:   return(0);
323: }

325: PETSC_INTERN PetscErrorCode PetscSFCreateLocalSF_Allgatherv(PetscSF sf,PetscSF *out)
326: {
327:   PetscInt       i,nroots,nleaves,rstart,*ilocal;
328:   PetscSFNode    *iremote;
329:   PetscSF        lsf;

333:   nleaves = sf->nleaves ? sf->nroots : 0; /* sf->nleaves can be zero with SFGather(v) */
334:   nroots  = nleaves;
335:   PetscMalloc1(nleaves,&ilocal);
336:   PetscMalloc1(nleaves,&iremote);
337:   PetscLayoutGetRange(sf->map,&rstart,NULL);

339:   for (i=0; i<nleaves; i++) {
340:     ilocal[i]        = rstart + i; /* lsf does not change leave indices */
341:     iremote[i].rank  = 0;          /* rank in PETSC_COMM_SELF */
342:     iremote[i].index = i;          /* root index */
343:   }

345:   PetscSFCreate(PETSC_COMM_SELF,&lsf);
346:   PetscSFSetGraph(lsf,nroots,nleaves,ilocal,PETSC_OWN_POINTER,iremote,PETSC_OWN_POINTER);
347:   PetscSFSetUp(lsf);
348:   *out = lsf;
349:   return(0);
350: }

352: PETSC_INTERN PetscErrorCode PetscSFCreate_Allgatherv(PetscSF sf)
353: {
354:   PetscErrorCode     ierr;
355:   PetscSF_Allgatherv *dat = (PetscSF_Allgatherv*)sf->data;

358:   sf->ops->BcastAndOpEnd   = PetscSFBcastAndOpEnd_Basic;
359:   sf->ops->ReduceEnd       = PetscSFReduceEnd_Basic;

361:   sf->ops->SetUp           = PetscSFSetUp_Allgatherv;
362:   sf->ops->Reset           = PetscSFReset_Allgatherv;
363:   sf->ops->Destroy         = PetscSFDestroy_Allgatherv;
364:   sf->ops->GetRootRanks    = PetscSFGetRootRanks_Allgatherv;
365:   sf->ops->GetLeafRanks    = PetscSFGetLeafRanks_Allgatherv;
366:   sf->ops->GetGraph        = PetscSFGetGraph_Allgatherv;
367:   sf->ops->BcastAndOpBegin = PetscSFBcastAndOpBegin_Allgatherv;
368:   sf->ops->ReduceBegin     = PetscSFReduceBegin_Allgatherv;
369:   sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Allgatherv;
370:   sf->ops->FetchAndOpEnd   = PetscSFFetchAndOpEnd_Allgatherv;
371:   sf->ops->CreateLocalSF   = PetscSFCreateLocalSF_Allgatherv;
372:   sf->ops->BcastToZero     = PetscSFBcastToZero_Allgatherv;

374:   PetscNewLog(sf,&dat);
375:   sf->data = (void*)dat;
376:   return(0);
377: }