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