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