Actual source code: sfallgatherv.c
1: #include <../src/vec/is/sf/impls/basic/allgatherv/sfallgatherv.h>
3: PETSC_INTERN PetscErrorCode PetscSFBcastBegin_Gatherv(PetscSF, MPI_Datatype, PetscMemType, const void *, PetscMemType, void *, MPI_Op);
5: /* PetscSFGetGraph is non-collective. An implementation should not have collective calls */
6: PETSC_INTERN PetscErrorCode PetscSFGetGraph_Allgatherv(PetscSF sf, PetscInt *nroots, PetscInt *nleaves, const PetscInt **ilocal, const PetscSFNode **iremote)
7: {
8: PetscInt i, j, k;
9: const PetscInt *range;
10: PetscMPIInt size;
12: MPI_Comm_size(PetscObjectComm((PetscObject)sf), &size);
13: if (nroots) *nroots = sf->nroots;
14: if (nleaves) *nleaves = sf->nleaves;
15: if (ilocal) *ilocal = NULL; /* Contiguous leaves */
16: if (iremote) {
17: if (!sf->remote && sf->nleaves) { /* The && sf->nleaves makes sfgatherv able to inherit this routine */
18: PetscLayoutGetRanges(sf->map, &range);
19: PetscMalloc1(sf->nleaves, &sf->remote);
20: sf->remote_alloc = sf->remote;
21: for (i = 0; i < size; i++) {
22: for (j = range[i], k = 0; j < range[i + 1]; j++, k++) {
23: sf->remote[j].rank = i;
24: sf->remote[j].index = k;
25: }
26: }
27: }
28: *iremote = sf->remote;
29: }
30: return 0;
31: }
33: PETSC_INTERN PetscErrorCode PetscSFSetUp_Allgatherv(PetscSF sf)
34: {
35: PetscSF_Allgatherv *dat = (PetscSF_Allgatherv *)sf->data;
36: PetscMPIInt size;
37: PetscInt i;
38: const PetscInt *range;
39: MPI_Comm comm;
41: PetscSFSetUp_Allgather(sf);
42: PetscObjectGetComm((PetscObject)sf, &comm);
43: MPI_Comm_size(comm, &size);
44: if (sf->nleaves) { /* This if (sf->nleaves) test makes sfgatherv able to inherit this routine */
45: PetscBool isallgatherv = PETSC_FALSE;
47: PetscMalloc1(size, &dat->recvcounts);
48: PetscMalloc1(size, &dat->displs);
49: PetscLayoutGetRanges(sf->map, &range);
51: for (i = 0; i < size; i++) {
52: PetscMPIIntCast(range[i], &dat->displs[i]);
53: PetscMPIIntCast(range[i + 1] - range[i], &dat->recvcounts[i]);
54: }
56: /* check if we actually have a one-to-all pattern */
57: PetscObjectTypeCompare((PetscObject)sf, PETSCSFALLGATHERV, &isallgatherv);
58: if (isallgatherv) {
59: PetscMPIInt rank, nRanksWithZeroRoots;
61: nRanksWithZeroRoots = (sf->nroots == 0) ? 1 : 0; /* I have no roots */
62: MPIU_Allreduce(MPI_IN_PLACE, &nRanksWithZeroRoots, 1, MPI_INT, MPI_SUM, comm);
63: if (nRanksWithZeroRoots == size - 1) { /* Only one rank has roots, which indicates a bcast pattern */
64: dat->bcast_pattern = PETSC_TRUE;
65: MPI_Comm_rank(comm, &rank);
66: dat->bcast_root = sf->nroots > 0 ? rank : -1;
67: MPIU_Allreduce(MPI_IN_PLACE, &dat->bcast_root, 1, MPI_INT, MPI_MAX, comm);
68: }
69: }
70: }
71: return 0;
72: }
74: PETSC_INTERN PetscErrorCode PetscSFReset_Allgatherv(PetscSF sf)
75: {
76: PetscSF_Allgatherv *dat = (PetscSF_Allgatherv *)sf->data;
77: PetscSFLink link = dat->avail, next;
79: PetscFree(dat->iranks);
80: PetscFree(dat->ioffset);
81: PetscFree(dat->irootloc);
82: PetscFree(dat->recvcounts);
83: PetscFree(dat->displs);
85: for (; link; link = next) {
86: next = link->next;
87: PetscSFLinkDestroy(sf, link);
88: }
89: dat->avail = NULL;
90: return 0;
91: }
93: PETSC_INTERN PetscErrorCode PetscSFDestroy_Allgatherv(PetscSF sf)
94: {
95: PetscSFReset_Allgatherv(sf);
96: PetscFree(sf->data);
97: return 0;
98: }
100: static PetscErrorCode PetscSFBcastBegin_Allgatherv(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, const void *rootdata, PetscMemType leafmtype, void *leafdata, MPI_Op op)
101: {
102: PetscSFLink link;
103: PetscMPIInt sendcount, rank;
104: MPI_Comm comm;
105: void *rootbuf = NULL, *leafbuf = NULL;
106: MPI_Request *req;
107: PetscSF_Allgatherv *dat = (PetscSF_Allgatherv *)sf->data;
109: PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_BCAST, &link);
110: PetscSFLinkPackRootData(sf, link, PETSCSF_REMOTE, rootdata);
111: PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */);
112: PetscObjectGetComm((PetscObject)sf, &comm);
113: MPI_Comm_rank(comm, &rank);
114: PetscMPIIntCast(sf->nroots, &sendcount);
115: PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_../../../../../..2LEAF, &rootbuf, &leafbuf, &req, NULL);
117: if (dat->bcast_pattern && rank == dat->bcast_root) (*link->Memcpy)(link, link->leafmtype_mpi, leafbuf, link->rootmtype_mpi, rootbuf, (size_t)sendcount * link->unitbytes);
118: /* Ready the buffers for MPI */
119: PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_../../../../../..2LEAF);
120: if (dat->bcast_pattern) MPIU_Ibcast(leafbuf, sf->nleaves, unit, dat->bcast_root, comm, req);
121: else MPIU_Iallgatherv(rootbuf, sendcount, unit, leafbuf, dat->recvcounts, dat->displs, unit, comm, req);
122: return 0;
123: }
125: static PetscErrorCode PetscSFReduceBegin_Allgatherv(PetscSF sf, MPI_Datatype unit, PetscMemType leafmtype, const void *leafdata, PetscMemType rootmtype, void *rootdata, MPI_Op op)
126: {
127: PetscSFLink link;
128: PetscSF_Allgatherv *dat = (PetscSF_Allgatherv *)sf->data;
129: PetscInt rstart;
130: PetscMPIInt rank, count, recvcount;
131: MPI_Comm comm;
132: void *rootbuf = NULL, *leafbuf = NULL;
133: MPI_Request *req;
135: PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_REDUCE, &link);
136: if (op == MPI_REPLACE) {
137: /* REPLACE is only meaningful when all processes have the same leafdata to reduce. Therefore copying from local leafdata is fine */
138: PetscLayoutGetRange(sf->map, &rstart, NULL);
139: (*link->Memcpy)(link, rootmtype, rootdata, leafmtype, (const char *)leafdata + (size_t)rstart * link->unitbytes, (size_t)sf->nroots * link->unitbytes);
140: if (PetscMemTypeDevice(leafmtype) && PetscMemTypeHost(rootmtype)) (*link->SyncStream)(link);
141: } else {
142: PetscObjectGetComm((PetscObject)sf, &comm);
143: PetscSFLinkPackLeafData(sf, link, PETSCSF_REMOTE, leafdata);
144: PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */);
145: PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_LEAF2../../../../../.., &rootbuf, &leafbuf, &req, NULL);
146: PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_LEAF2../../../../../..);
147: if (dat->bcast_pattern) {
148: #if defined(PETSC_HAVE_OMPI_MAJOR_VERSION) /* Workaround: cuda-aware OpenMPI-4.1.3 does not support MPI_Ireduce() with device buffers */
149: *req = MPI_REQUEST_NULL; /* Set NULL so that we can safely MPI_Wait(req) */
150: MPI_Reduce(leafbuf, rootbuf, sf->nleaves, unit, op, dat->bcast_root, comm);
151: #else
152: MPIU_Ireduce(leafbuf, rootbuf, sf->nleaves, unit, op, dat->bcast_root, comm, req);
153: #endif
154: } else { /* Reduce leafdata, then scatter to rootdata */
155: MPI_Comm_rank(comm, &rank);
156: PetscMPIIntCast(dat->rootbuflen[PETSCSF_REMOTE], &recvcount);
157: /* Allocate a separate leaf buffer on rank 0 */
158: if (rank == 0 && !link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi]) {
159: PetscSFMalloc(sf, link->leafmtype_mpi, sf->leafbuflen[PETSCSF_REMOTE] * link->unitbytes, (void **)&link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi]);
160: }
161: /* 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 */
162: if (rank == 0 && link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi] == leafbuf) leafbuf = MPI_IN_PLACE;
163: PetscMPIIntCast(sf->nleaves * link->bs, &count);
164: MPI_Reduce(leafbuf, link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi], count, link->basicunit, op, 0, comm); /* Must do reduce with MPI builtin datatype basicunit */
165: MPIU_Iscatterv(link->leafbuf_alloc[PETSCSF_REMOTE][link->leafmtype_mpi], dat->recvcounts, dat->displs, unit, rootbuf, recvcount, unit, 0, comm, req);
166: }
167: }
168: return 0;
169: }
171: PETSC_INTERN PetscErrorCode PetscSFReduceEnd_Allgatherv(PetscSF sf, MPI_Datatype unit, const void *leafdata, void *rootdata, MPI_Op op)
172: {
173: PetscSFLink link;
175: if (op == MPI_REPLACE) {
176: /* A rare case happens when op is MPI_REPLACE, using GPUs but no GPU aware MPI. In PetscSFReduceBegin_Allgather(v),
177: we did a device to device copy and in effect finished the communication. But in PetscSFLinkFinishCommunication()
178: of PetscSFReduceEnd_Basic(), it thinks since there is rootbuf, it calls PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI().
179: It does a host to device memory copy on rootbuf, wrongly overwriting the results. So we don't overload
180: PetscSFReduceEnd_Basic() in this case, and just reclaim the link.
181: */
182: PetscSFLinkGetInUse(sf, unit, rootdata, leafdata, PETSC_OWN_POINTER, &link);
183: PetscSFLinkReclaim(sf, &link);
184: } else {
185: PetscSFReduceEnd_Basic(sf, unit, leafdata, rootdata, op);
186: }
187: return 0;
188: }
190: static PetscErrorCode PetscSFBcastToZero_Allgatherv(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, const void *rootdata, PetscMemType leafmtype, void *leafdata)
191: {
192: PetscSFLink link;
193: PetscMPIInt rank;
195: PetscSFBcastBegin_Gatherv(sf, unit, rootmtype, rootdata, leafmtype, leafdata, MPI_REPLACE);
196: PetscSFLinkGetInUse(sf, unit, rootdata, leafdata, PETSC_OWN_POINTER, &link);
197: PetscSFLinkFinishCommunication(sf, link, PETSCSF_../../../../../..2LEAF);
198: MPI_Comm_rank(PetscObjectComm((PetscObject)sf), &rank);
199: if (rank == 0 && PetscMemTypeDevice(leafmtype) && !sf->use_gpu_aware_mpi) (*link->Memcpy)(link, PETSC_MEMTYPE_DEVICE, leafdata, PETSC_MEMTYPE_HOST, link->leafbuf[PETSC_MEMTYPE_HOST], sf->leafbuflen[PETSCSF_REMOTE] * link->unitbytes);
200: PetscSFLinkReclaim(sf, &link);
201: return 0;
202: }
204: /* 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).
206: 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
207: 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
208: 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
209: 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
210: 0,1,2 is 1,3,6 respectively. root is 10.
212: We use a simpler implementation. From the same initial state, we copy leafdata to leafupdate
213: rank-0 rank-1 rank-2
214: Root 1
215: Leaf 2 3 4
216: Leafupdate 2 3 4
218: Do MPI_Exscan on leafupdate,
219: rank-0 rank-1 rank-2
220: Root 1
221: Leaf 2 3 4
222: Leafupdate 2 2 5
224: BcastAndOp from root to leafupdate,
225: rank-0 rank-1 rank-2
226: Root 1
227: Leaf 2 3 4
228: Leafupdate 3 3 6
230: Copy root to leafupdate on rank-0
231: rank-0 rank-1 rank-2
232: Root 1
233: Leaf 2 3 4
234: Leafupdate 1 3 6
236: Reduce from leaf to root,
237: rank-0 rank-1 rank-2
238: Root 10
239: Leaf 2 3 4
240: Leafupdate 1 3 6
241: */
242: PETSC_INTERN PetscErrorCode PetscSFFetchAndOpBegin_Allgatherv(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, void *rootdata, PetscMemType leafmtype, const void *leafdata, void *leafupdate, MPI_Op op)
243: {
244: PetscSFLink link;
245: MPI_Comm comm;
246: PetscMPIInt count;
248: PetscObjectGetComm((PetscObject)sf, &comm);
250: /* Copy leafdata to leafupdate */
251: PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_FETCH, &link);
252: PetscSFLinkPackLeafData(sf, link, PETSCSF_REMOTE, leafdata); /* Sync the device */
253: (*link->Memcpy)(link, leafmtype, leafupdate, leafmtype, leafdata, sf->nleaves * link->unitbytes);
254: PetscSFLinkGetInUse(sf, unit, rootdata, leafdata, PETSC_OWN_POINTER, &link);
256: /* Exscan on leafupdate and then BcastAndOp rootdata to leafupdate */
257: if (op == MPI_REPLACE) {
258: PetscMPIInt size, rank, prev, next;
259: MPI_Comm_rank(comm, &rank);
260: MPI_Comm_size(comm, &size);
261: prev = rank ? rank - 1 : MPI_PROC_NULL;
262: next = (rank < size - 1) ? rank + 1 : MPI_PROC_NULL;
263: PetscMPIIntCast(sf->nleaves, &count);
264: MPI_Sendrecv_replace(leafupdate, count, unit, next, link->tag, prev, link->tag, comm, MPI_STATUSES_IGNORE);
265: } else {
266: PetscMPIIntCast(sf->nleaves * link->bs, &count);
267: MPI_Exscan(MPI_IN_PLACE, leafupdate, count, link->basicunit, op, comm);
268: }
269: PetscSFLinkReclaim(sf, &link);
270: PetscSFBcastBegin(sf, unit, rootdata, leafupdate, op);
271: PetscSFBcastEnd(sf, unit, rootdata, leafupdate, op);
273: /* Bcast roots to rank 0's leafupdate */
274: PetscSFBcastToZero_Private(sf, unit, rootdata, leafupdate); /* Using this line makes Allgather SFs able to inherit this routine */
276: /* Reduce leafdata to rootdata */
277: PetscSFReduceBegin(sf, unit, leafdata, rootdata, op);
278: return 0;
279: }
281: PETSC_INTERN PetscErrorCode PetscSFFetchAndOpEnd_Allgatherv(PetscSF sf, MPI_Datatype unit, void *rootdata, const void *leafdata, void *leafupdate, MPI_Op op)
282: {
283: PetscSFReduceEnd(sf, unit, leafdata, rootdata, op);
284: return 0;
285: }
287: /* Get root ranks accessing my leaves */
288: PETSC_INTERN PetscErrorCode PetscSFGetRootRanks_Allgatherv(PetscSF sf, PetscInt *nranks, const PetscMPIInt **ranks, const PetscInt **roffset, const PetscInt **rmine, const PetscInt **rremote)
289: {
290: PetscInt i, j, k, size;
291: const PetscInt *range;
293: /* Lazily construct these large arrays if users really need them for this type of SF. Very likely, they do not */
294: if (sf->nranks && !sf->ranks) { /* On rank!=0, sf->nranks=0. The sf->nranks test makes this routine also works for sfgatherv */
295: size = sf->nranks;
296: PetscLayoutGetRanges(sf->map, &range);
297: PetscMalloc4(size, &sf->ranks, size + 1, &sf->roffset, sf->nleaves, &sf->rmine, sf->nleaves, &sf->rremote);
298: for (i = 0; i < size; i++) sf->ranks[i] = i;
299: PetscArraycpy(sf->roffset, range, size + 1);
300: for (i = 0; i < sf->nleaves; i++) sf->rmine[i] = i; /*rmine are never NULL even for contiguous leaves */
301: for (i = 0; i < size; i++) {
302: for (j = range[i], k = 0; j < range[i + 1]; j++, k++) sf->rremote[j] = k;
303: }
304: }
306: if (nranks) *nranks = sf->nranks;
307: if (ranks) *ranks = sf->ranks;
308: if (roffset) *roffset = sf->roffset;
309: if (rmine) *rmine = sf->rmine;
310: if (rremote) *rremote = sf->rremote;
311: return 0;
312: }
314: /* Get leaf ranks accessing my roots */
315: PETSC_INTERN PetscErrorCode PetscSFGetLeafRanks_Allgatherv(PetscSF sf, PetscInt *niranks, const PetscMPIInt **iranks, const PetscInt **ioffset, const PetscInt **irootloc)
316: {
317: PetscSF_Allgatherv *dat = (PetscSF_Allgatherv *)sf->data;
318: MPI_Comm comm;
319: PetscMPIInt size, rank;
320: PetscInt i, j;
322: /* Lazily construct these large arrays if users really need them for this type of SF. Very likely, they do not */
323: PetscObjectGetComm((PetscObject)sf, &comm);
324: MPI_Comm_size(comm, &size);
325: MPI_Comm_rank(comm, &rank);
326: if (niranks) *niranks = size;
328: /* PetscSF_Basic has distinguished incoming ranks. Here we do not need that. But we must put self as the first and
329: sort other ranks. See comments in PetscSFSetUp_Basic about MatGetBrowsOfAoCols_MPIAIJ on why.
330: */
331: if (iranks) {
332: if (!dat->iranks) {
333: PetscMalloc1(size, &dat->iranks);
334: dat->iranks[0] = rank;
335: for (i = 0, j = 1; i < size; i++) {
336: if (i == rank) continue;
337: dat->iranks[j++] = i;
338: }
339: }
340: *iranks = dat->iranks; /* dat->iranks was init'ed to NULL by PetscNew */
341: }
343: if (ioffset) {
344: if (!dat->ioffset) {
345: PetscMalloc1(size + 1, &dat->ioffset);
346: for (i = 0; i <= size; i++) dat->ioffset[i] = i * sf->nroots;
347: }
348: *ioffset = dat->ioffset;
349: }
351: if (irootloc) {
352: if (!dat->irootloc) {
353: PetscMalloc1(sf->nleaves, &dat->irootloc);
354: for (i = 0; i < size; i++) {
355: for (j = 0; j < sf->nroots; j++) dat->irootloc[i * sf->nroots + j] = j;
356: }
357: }
358: *irootloc = dat->irootloc;
359: }
360: return 0;
361: }
363: PETSC_INTERN PetscErrorCode PetscSFCreateLocalSF_Allgatherv(PetscSF sf, PetscSF *out)
364: {
365: PetscInt i, nroots, nleaves, rstart, *ilocal;
366: PetscSFNode *iremote;
367: PetscSF lsf;
369: nleaves = sf->nleaves ? sf->nroots : 0; /* sf->nleaves can be zero with SFGather(v) */
370: nroots = nleaves;
371: PetscMalloc1(nleaves, &ilocal);
372: PetscMalloc1(nleaves, &iremote);
373: PetscLayoutGetRange(sf->map, &rstart, NULL);
375: for (i = 0; i < nleaves; i++) {
376: ilocal[i] = rstart + i; /* lsf does not change leave indices */
377: iremote[i].rank = 0; /* rank in PETSC_COMM_SELF */
378: iremote[i].index = i; /* root index */
379: }
381: PetscSFCreate(PETSC_COMM_SELF, &lsf);
382: PetscSFSetGraph(lsf, nroots, nleaves, ilocal, PETSC_OWN_POINTER, iremote, PETSC_OWN_POINTER);
383: PetscSFSetUp(lsf);
384: *out = lsf;
385: return 0;
386: }
388: PETSC_INTERN PetscErrorCode PetscSFCreate_Allgatherv(PetscSF sf)
389: {
390: PetscSF_Allgatherv *dat = (PetscSF_Allgatherv *)sf->data;
392: sf->ops->BcastEnd = PetscSFBcastEnd_Basic;
393: sf->ops->ReduceEnd = PetscSFReduceEnd_Allgatherv;
395: sf->ops->SetUp = PetscSFSetUp_Allgatherv;
396: sf->ops->Reset = PetscSFReset_Allgatherv;
397: sf->ops->Destroy = PetscSFDestroy_Allgatherv;
398: sf->ops->GetRootRanks = PetscSFGetRootRanks_Allgatherv;
399: sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Allgatherv;
400: sf->ops->GetGraph = PetscSFGetGraph_Allgatherv;
401: sf->ops->BcastBegin = PetscSFBcastBegin_Allgatherv;
402: sf->ops->ReduceBegin = PetscSFReduceBegin_Allgatherv;
403: sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Allgatherv;
404: sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv;
405: sf->ops->CreateLocalSF = PetscSFCreateLocalSF_Allgatherv;
406: sf->ops->BcastToZero = PetscSFBcastToZero_Allgatherv;
408: PetscNew(&dat);
409: dat->bcast_root = -1;
410: sf->data = (void *)dat;
411: return 0;
412: }