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