Actual source code: sfalltoall.c

petsc-master 2019-11-16
Report Typos and Errors
  1:  #include <../src/vec/is/sf/impls/basic/allgatherv/sfallgatherv.h>
  2:  #include <../src/vec/is/sf/impls/basic/allgather/sfallgather.h>
  3:  #include <../src/vec/is/sf/impls/basic/gatherv/sfgatherv.h>

  5: #define PetscSFPackGet_Alltoall PetscSFPackGet_Allgatherv

  7: /* Reuse the type. The difference is some fields (i.e., displs, recvcounts) are not used, which is not a big deal */
  8: typedef PetscSF_Allgatherv PetscSF_Alltoall;

 10: /*===================================================================================*/
 11: /*              Implementations of SF public APIs                                    */
 12: /*===================================================================================*/
 13: static PetscErrorCode PetscSFGetGraph_Alltoall(PetscSF sf,PetscInt *nroots,PetscInt *nleaves,const PetscInt **ilocal,const PetscSFNode **iremote)
 14: {
 16:   PetscInt       i;

 19:   if (nroots)  *nroots  = sf->nroots;
 20:   if (nleaves) *nleaves = sf->nleaves;
 21:   if (ilocal)  *ilocal  = NULL; /* Contiguous local indices */
 22:   if (iremote) {
 23:     if (!sf->remote) {
 24:       PetscMalloc1(sf->nleaves,&sf->remote);
 25:       sf->remote_alloc = sf->remote;
 26:       for (i=0; i<sf->nleaves; i++) {
 27:         sf->remote[i].rank  = i;
 28:         sf->remote[i].index = i;
 29:       }
 30:     }
 31:     *iremote = sf->remote;
 32:   }
 33:   return(0);
 34: }

 36: static PetscErrorCode PetscSFBcastAndOpBegin_Alltoall(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op)
 37: {
 38:   PetscErrorCode       ierr;
 39:   PetscSFPack          link;
 40:   MPI_Comm             comm;
 41:   char                 *recvbuf;

 44:   PetscSFPackGet_Alltoall(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);
 45:   PetscObjectGetComm((PetscObject)sf,&comm);

 47:   if (op != MPIU_REPLACE) {
 48:     if (!link->leafbuf[leafmtype]) {PetscMallocWithMemType(leafmtype,sf->nleaves*link->unitbytes,(void**)&link->leafbuf[leafmtype]);}
 49:     recvbuf = link->leafbuf[leafmtype];
 50:   } else {
 51:     recvbuf = (char*)leafdata;
 52:   }
 53:   MPIU_Ialltoall(rootdata,1,unit,recvbuf,1,unit,comm,link->rootreqs[PETSCSF_../../../../../..2LEAF_BCAST][rootmtype]);
 54:   return(0);
 55: }

 57: static PetscErrorCode PetscSFReduceBegin_Alltoall(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op)
 58: {
 59:   PetscErrorCode       ierr;
 60:   PetscSFPack          link;
 61:   MPI_Comm             comm;
 62:   char                 *recvbuf;

 65:   PetscSFPackGet_Alltoall(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);
 66:   PetscObjectGetComm((PetscObject)sf,&comm);

 68:   if (op != MPIU_REPLACE) {
 69:     if (!link->rootbuf[rootmtype]) {PetscMallocWithMemType(rootmtype,sf->nroots*link->unitbytes,(void**)&link->rootbuf[rootmtype]);}
 70:     recvbuf = link->rootbuf[rootmtype];
 71:   } else {
 72:     recvbuf = (char*)rootdata;
 73:   }
 74:   MPIU_Ialltoall(leafdata,1,unit,recvbuf,1,unit,comm,link->rootreqs[PETSCSF_LEAF2../../../../../.._REDUCE][rootmtype]);
 75:   return(0);
 76: }

 78: static PetscErrorCode PetscSFCreateLocalSF_Alltoall(PetscSF sf,PetscSF *out)
 79: {
 81:   PetscInt       nroots=1,nleaves=1,*ilocal;
 82:   PetscSFNode    *iremote = NULL;
 83:   PetscSF        lsf;
 84:   PetscMPIInt    rank;

 87:   nroots  = 1;
 88:   nleaves = 1;
 89:   MPI_Comm_rank(PetscObjectComm((PetscObject)sf),&rank);
 90:   PetscMalloc1(nleaves,&ilocal);
 91:   PetscMalloc1(nleaves,&iremote);
 92:   ilocal[0]        = rank;
 93:   iremote[0].rank  = 0;    /* rank in PETSC_COMM_SELF */
 94:   iremote[0].index = rank; /* LocalSF is an embedded SF. Indices are not remapped */

 96:   PetscSFCreate(PETSC_COMM_SELF,&lsf);
 97:   PetscSFSetGraph(lsf,nroots,nleaves,NULL/*contiguous leaves*/,PETSC_OWN_POINTER,iremote,PETSC_OWN_POINTER);
 98:   PetscSFSetUp(lsf);
 99:   *out = lsf;
100:   return(0);
101: }

103: static PetscErrorCode PetscSFCreateEmbeddedSF_Alltoall(PetscSF sf,PetscInt nselected,const PetscInt *selected,PetscSF *newsf)
104: {
106:   PetscInt       i,*tmproots,*ilocal,ndranks,ndiranks;
107:   PetscSFNode    *iremote;
108:   PetscMPIInt    nroots,*roots,nleaves,*leaves,rank;
109:   MPI_Comm       comm;
110:   PetscSF_Basic  *bas;
111:   PetscSF        esf;

114:   PetscObjectGetComm((PetscObject)sf,&comm);
115:   MPI_Comm_rank(comm,&rank);

117:   /* Uniq selected[] and store the result in roots[] */
118:   PetscMalloc1(nselected,&tmproots);
119:   PetscArraycpy(tmproots,selected,nselected);
120:   PetscSortRemoveDupsInt(&nselected,tmproots); /* nselected might be changed */
121:   if (tmproots[0] < 0 || tmproots[nselected-1] >= sf->nroots) SETERRQ3(comm,PETSC_ERR_ARG_OUTOFRANGE,"Min/Max root indices %D/%D are not in [0,%D)",tmproots[0],tmproots[nselected-1],sf->nroots);
122:   nroots = nselected;   /* For Alltoall, we know root indices will not overflow MPI_INT */
123:   PetscMalloc1(nselected,&roots);
124:   for (i=0; i<nselected; i++) roots[i] = tmproots[i];
125:   PetscFree(tmproots);

127:   /* Find out which leaves are still connected to roots in the embedded sf. Expect PetscCommBuildTwoSided is more scalable than MPI_Alltoall */
128:   PetscCommBuildTwoSided(comm,0/*empty msg*/,MPI_INT/*fake*/,nroots,roots,NULL/*todata*/,&nleaves,&leaves,NULL/*fromdata*/);

130:   /* Move myself ahead if rank is in leaves[], since I am a distinguished rank */
131:   ndranks = 0;
132:   for (i=0; i<nleaves; i++) {
133:     if (leaves[i] == rank) {leaves[i] = -rank; ndranks = 1; break;}
134:   }
135:   PetscSortMPIInt(nleaves,leaves);
136:   if (nleaves && leaves[0] < 0) leaves[0] = rank;

138:   /* Build esf and fill its fields manually (without calling PetscSFSetUp) */
139:   PetscMalloc1(nleaves,&ilocal);
140:   PetscMalloc1(nleaves,&iremote);
141:   for (i=0; i<nleaves; i++) { /* 1:1 map from roots to leaves */
142:     ilocal[i]        = leaves[i];
143:     iremote[i].rank  = leaves[i];
144:     iremote[i].index = leaves[i];
145:   }
146:   PetscSFCreate(comm,&esf);
147:   PetscSFSetType(esf,PETSCSFBASIC); /* This optimized routine can only create a basic sf */
148:   PetscSFSetGraph(esf,sf->nleaves,nleaves,ilocal,PETSC_OWN_POINTER,iremote,PETSC_OWN_POINTER);

150:   /* As if we are calling PetscSFSetUpRanks(esf,self's group) */
151:   PetscMalloc4(nleaves,&esf->ranks,nleaves+1,&esf->roffset,nleaves,&esf->rmine,nleaves,&esf->rremote);
152:   esf->nranks     = nleaves;
153:   esf->ndranks    = ndranks;
154:   esf->roffset[0] = 0;
155:   for (i=0; i<nleaves; i++) {
156:     esf->ranks[i]     = leaves[i];
157:     esf->roffset[i+1] = i+1;
158:     esf->rmine[i]     = leaves[i];
159:     esf->rremote[i]   = leaves[i];
160:   }

162:   /* Set up esf->data, the incoming communication (i.e., recv info), which is usually done by PetscSFSetUp_Basic */
163:   bas  = (PetscSF_Basic*)esf->data;
164:   PetscMalloc2(nroots,&bas->iranks,nroots+1,&bas->ioffset);
165:   PetscMalloc1(nroots,&bas->irootloc);
166:   /* Move myself ahead if rank is in roots[], since I am a distinguished irank */
167:   ndiranks = 0;
168:   for (i=0; i<nroots; i++) {
169:     if (roots[i] == rank) {roots[i] = -rank; ndiranks = 1; break;}
170:   }
171:   PetscSortMPIInt(nroots,roots);
172:   if (nroots && roots[0] < 0) roots[0] = rank;

174:   bas->niranks    = nroots;
175:   bas->ndiranks   = ndiranks;
176:   bas->ioffset[0] = 0;
177:   bas->itotal     = nroots;
178:   for (i=0; i<nroots; i++) {
179:     bas->iranks[i]    = roots[i];
180:     bas->ioffset[i+1] = i+1;
181:     bas->irootloc[i]  = roots[i];
182:   }

184:   *newsf = esf;
185:   return(0);
186: }

188: PETSC_INTERN PetscErrorCode PetscSFCreate_Alltoall(PetscSF sf)
189: {
190:   PetscErrorCode   ierr;
191:   PetscSF_Alltoall *dat = (PetscSF_Alltoall*)sf->data;

194:   /* Inherit from Allgatherv. It is astonishing Alltoall can inherit so much from Allgather(v) */
195:   sf->ops->Destroy         = PetscSFDestroy_Allgatherv;
196:   sf->ops->Reset           = PetscSFReset_Allgatherv;
197:   sf->ops->BcastAndOpEnd   = PetscSFBcastAndOpEnd_Allgatherv;
198:   sf->ops->ReduceEnd       = PetscSFReduceEnd_Allgatherv;
199:   sf->ops->FetchAndOpEnd   = PetscSFFetchAndOpEnd_Allgatherv;
200:   sf->ops->GetRootRanks    = PetscSFGetRootRanks_Allgatherv;

202:   /* Inherit from Allgather. Every process gathers equal-sized data from others, which enables this inheritance. */
203:   sf->ops->GetLeafRanks    = PetscSFGetLeafRanks_Allgatherv;

205:   /* Inherit from Gatherv. Each root has only one leaf connected, which enables this inheritance */
206:   sf->ops->FetchAndOpBegin  = PetscSFFetchAndOpBegin_Gatherv;

208:   /* Alltoall stuff */
209:   sf->ops->GetGraph         = PetscSFGetGraph_Alltoall;
210:   sf->ops->BcastAndOpBegin  = PetscSFBcastAndOpBegin_Alltoall;
211:   sf->ops->ReduceBegin      = PetscSFReduceBegin_Alltoall;
212:   sf->ops->CreateLocalSF    = PetscSFCreateLocalSF_Alltoall;
213:   sf->ops->CreateEmbeddedSF = PetscSFCreateEmbeddedSF_Alltoall;

215:   PetscNewLog(sf,&dat);
216:   sf->data = (void*)dat;
217:   return(0);
218: }