Actual source code: sfalltoall.c

  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: /* Reuse the type. The difference is some fields (i.e., displs, recvcounts) are not used, which is not a big deal */
  6: typedef PetscSF_Allgatherv PetscSF_Alltoall;

  8: static PetscErrorCode PetscSFLinkStartCommunication_Alltoall(PetscSF sf, PetscSFLink link, PetscSFDirection direction)
  9: {
 10:   MPI_Comm     comm    = MPI_COMM_NULL;
 11:   void        *rootbuf = NULL, *leafbuf = NULL;
 12:   MPI_Request *req  = NULL;
 13:   MPI_Datatype unit = link->unit;

 15:   PetscFunctionBegin;
 16:   if (direction == PETSCSF_ROOT2LEAF) {
 17:     PetscCall(PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */));
 18:   } else {
 19:     PetscCall(PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host */));
 20:   }
 21:   PetscCall(PetscObjectGetComm((PetscObject)sf, &comm));
 22:   PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, direction, &rootbuf, &leafbuf, &req, NULL));
 23:   PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link, direction));

 25:   if (direction == PETSCSF_ROOT2LEAF) {
 26:     PetscCallMPI(MPIU_Ialltoall(rootbuf, 1, unit, leafbuf, 1, unit, comm, req));
 27:   } else {
 28:     PetscCallMPI(MPIU_Ialltoall(leafbuf, 1, unit, rootbuf, 1, unit, comm, req));
 29:   }
 30:   PetscFunctionReturn(PETSC_SUCCESS);
 31: }

 33: static PetscErrorCode PetscSFSetCommunicationOps_Alltoall(PetscSF sf, PetscSFLink link)
 34: {
 35:   PetscFunctionBegin;
 36:   link->StartCommunication = PetscSFLinkStartCommunication_Alltoall;
 37:   PetscFunctionReturn(PETSC_SUCCESS);
 38: }

 40: /*===================================================================================*/
 41: /*              Implementations of SF public APIs                                    */
 42: /*===================================================================================*/
 43: static PetscErrorCode PetscSFGetGraph_Alltoall(PetscSF sf, PetscInt *nroots, PetscInt *nleaves, const PetscInt **ilocal, const PetscSFNode **iremote)
 44: {
 45:   PetscInt i;

 47:   PetscFunctionBegin;
 48:   if (nroots) *nroots = sf->nroots;
 49:   if (nleaves) *nleaves = sf->nleaves;
 50:   if (ilocal) *ilocal = NULL; /* Contiguous local indices */
 51:   if (iremote) {
 52:     if (!sf->remote) {
 53:       PetscCall(PetscMalloc1(sf->nleaves, &sf->remote));
 54:       sf->remote_alloc = sf->remote;
 55:       for (i = 0; i < sf->nleaves; i++) {
 56:         sf->remote[i].rank  = i;
 57:         sf->remote[i].index = i;
 58:       }
 59:     }
 60:     *iremote = sf->remote;
 61:   }
 62:   PetscFunctionReturn(PETSC_SUCCESS);
 63: }

 65: static PetscErrorCode PetscSFCreateLocalSF_Alltoall(PetscSF sf, PetscSF *out)
 66: {
 67:   PetscInt     nroots = 1, nleaves = 1, *ilocal;
 68:   PetscSFNode *iremote = NULL;
 69:   PetscSF      lsf;
 70:   PetscMPIInt  rank;

 72:   PetscFunctionBegin;
 73:   nroots  = 1;
 74:   nleaves = 1;
 75:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)sf), &rank));
 76:   PetscCall(PetscMalloc1(nleaves, &ilocal));
 77:   PetscCall(PetscMalloc1(nleaves, &iremote));
 78:   ilocal[0]        = rank;
 79:   iremote[0].rank  = 0;    /* rank in PETSC_COMM_SELF */
 80:   iremote[0].index = rank; /* LocalSF is an embedded SF. Indices are not remapped */

 82:   PetscCall(PetscSFCreate(PETSC_COMM_SELF, &lsf));
 83:   PetscCall(PetscSFSetGraph(lsf, nroots, nleaves, NULL /*contiguous leaves*/, PETSC_OWN_POINTER, iremote, PETSC_OWN_POINTER));
 84:   PetscCall(PetscSFSetUp(lsf));
 85:   *out = lsf;
 86:   PetscFunctionReturn(PETSC_SUCCESS);
 87: }

 89: static PetscErrorCode PetscSFCreateEmbeddedRootSF_Alltoall(PetscSF sf, PetscInt nselected, const PetscInt *selected, PetscSF *newsf)
 90: {
 91:   PetscInt       i, *tmproots, *ilocal, ndranks, ndiranks;
 92:   PetscSFNode   *iremote;
 93:   PetscMPIInt    nroots, *roots, nleaves, *leaves, rank;
 94:   MPI_Comm       comm;
 95:   PetscSF_Basic *bas;
 96:   PetscSF        esf;

 98:   PetscFunctionBegin;
 99:   PetscCall(PetscObjectGetComm((PetscObject)sf, &comm));
100:   PetscCallMPI(MPI_Comm_rank(comm, &rank));

102:   /* Uniq selected[] and store the result in roots[] */
103:   PetscCall(PetscMalloc1(nselected, &tmproots));
104:   PetscCall(PetscArraycpy(tmproots, selected, nselected));
105:   PetscCall(PetscSortRemoveDupsInt(&nselected, tmproots)); /* nselected might be changed */
106:   PetscCheck(tmproots[0] >= 0 && tmproots[nselected - 1] < sf->nroots, comm, PETSC_ERR_ARG_OUTOFRANGE, "Min/Max root indices %" PetscInt_FMT "/%" PetscInt_FMT " are not in [0,%" PetscInt_FMT ")", tmproots[0], tmproots[nselected - 1], sf->nroots);
107:   nroots = nselected; /* For Alltoall, we know root indices will not overflow MPI_INT */
108:   PetscCall(PetscMalloc1(nselected, &roots));
109:   for (i = 0; i < nselected; i++) roots[i] = tmproots[i];
110:   PetscCall(PetscFree(tmproots));

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

115:   /* Move myself ahead if rank is in leaves[], since I am a distinguished rank */
116:   ndranks = 0;
117:   for (i = 0; i < nleaves; i++) {
118:     if (leaves[i] == rank) {
119:       leaves[i] = -rank;
120:       ndranks   = 1;
121:       break;
122:     }
123:   }
124:   PetscCall(PetscSortMPIInt(nleaves, leaves));
125:   if (nleaves && leaves[0] < 0) leaves[0] = rank;

127:   /* Build esf and fill its fields manually (without calling PetscSFSetUp) */
128:   PetscCall(PetscMalloc1(nleaves, &ilocal));
129:   PetscCall(PetscMalloc1(nleaves, &iremote));
130:   for (i = 0; i < nleaves; i++) { /* 1:1 map from roots to leaves */
131:     ilocal[i]        = leaves[i];
132:     iremote[i].rank  = leaves[i];
133:     iremote[i].index = leaves[i];
134:   }
135:   PetscCall(PetscSFCreate(comm, &esf));
136:   PetscCall(PetscSFSetType(esf, PETSCSFBASIC)); /* This optimized routine can only create a basic sf */
137:   PetscCall(PetscSFSetGraph(esf, sf->nleaves, nleaves, ilocal, PETSC_OWN_POINTER, iremote, PETSC_OWN_POINTER));

139:   /* As if we are calling PetscSFSetUpRanks(esf,self's group) */
140:   PetscCall(PetscMalloc4(nleaves, &esf->ranks, nleaves + 1, &esf->roffset, nleaves, &esf->rmine, nleaves, &esf->rremote));
141:   esf->nranks     = nleaves;
142:   esf->ndranks    = ndranks;
143:   esf->roffset[0] = 0;
144:   for (i = 0; i < nleaves; i++) {
145:     esf->ranks[i]       = leaves[i];
146:     esf->roffset[i + 1] = i + 1;
147:     esf->rmine[i]       = leaves[i];
148:     esf->rremote[i]     = leaves[i];
149:   }

151:   /* Set up esf->data, the incoming communication (i.e., recv info), which is usually done by PetscSFSetUp_Basic */
152:   bas = (PetscSF_Basic *)esf->data;
153:   PetscCall(PetscMalloc2(nroots, &bas->iranks, nroots + 1, &bas->ioffset));
154:   PetscCall(PetscMalloc1(nroots, &bas->irootloc));
155:   /* Move myself ahead if rank is in roots[], since I am a distinguished irank */
156:   ndiranks = 0;
157:   for (i = 0; i < nroots; i++) {
158:     if (roots[i] == rank) {
159:       roots[i] = -rank;
160:       ndiranks = 1;
161:       break;
162:     }
163:   }
164:   PetscCall(PetscSortMPIInt(nroots, roots));
165:   if (nroots && roots[0] < 0) roots[0] = rank;

167:   bas->niranks    = nroots;
168:   bas->ndiranks   = ndiranks;
169:   bas->ioffset[0] = 0;
170:   bas->itotal     = nroots;
171:   for (i = 0; i < nroots; i++) {
172:     bas->iranks[i]      = roots[i];
173:     bas->ioffset[i + 1] = i + 1;
174:     bas->irootloc[i]    = roots[i];
175:   }

177:   /* See PetscSFCreateEmbeddedRootSF_Basic */
178:   esf->nleafreqs  = esf->nranks - esf->ndranks;
179:   bas->nrootreqs  = bas->niranks - bas->ndiranks;
180:   esf->persistent = PETSC_TRUE;
181:   /* Setup packing related fields */
182:   PetscCall(PetscSFSetUpPackFields(esf));

184:   esf->setupcalled = PETSC_TRUE; /* We have done setup ourselves! */
185:   *newsf           = esf;
186:   PetscFunctionReturn(PETSC_SUCCESS);
187: }

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

193:   PetscFunctionBegin;
194:   sf->ops->BcastBegin  = PetscSFBcastBegin_Basic;
195:   sf->ops->BcastEnd    = PetscSFBcastEnd_Basic;
196:   sf->ops->ReduceBegin = PetscSFReduceBegin_Basic;
197:   sf->ops->ReduceEnd   = PetscSFReduceEnd_Basic;

199:   /* Inherit from Allgatherv. It is astonishing Alltoall can inherit so much from Allgather(v) */
200:   sf->ops->Destroy       = PetscSFDestroy_Allgatherv;
201:   sf->ops->Reset         = PetscSFReset_Allgatherv;
202:   sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv;
203:   sf->ops->GetRootRanks  = PetscSFGetRootRanks_Allgatherv;

205:   /* Inherit from Allgather. Every process gathers equal-sized data from others, which enables this inheritance. */
206:   sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Allgatherv;
207:   sf->ops->SetUp        = PetscSFSetUp_Allgather;

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

212:   /* Alltoall stuff */
213:   sf->ops->GetGraph             = PetscSFGetGraph_Alltoall;
214:   sf->ops->CreateLocalSF        = PetscSFCreateLocalSF_Alltoall;
215:   sf->ops->CreateEmbeddedRootSF = PetscSFCreateEmbeddedRootSF_Alltoall;

217:   sf->ops->SetCommunicationOps = PetscSFSetCommunicationOps_Alltoall;

219:   sf->collective = PETSC_TRUE;

221:   PetscCall(PetscNew(&dat));
222:   sf->data = (void *)dat;
223:   PetscFunctionReturn(PETSC_SUCCESS);
224: }