Actual source code: sfbasic.c
1: #include <../src/vec/is/sf/impls/basic/sfbasic.h>
2: #include <../src/vec/is/sf/impls/basic/sfpack.h>
3: #include <petsc/private/viewerimpl.h>
5: /*===================================================================================*/
6: /* SF public interface implementations */
7: /*===================================================================================*/
8: PETSC_INTERN PetscErrorCode PetscSFSetUp_Basic(PetscSF sf)
9: {
10: PetscSF_Basic *bas = (PetscSF_Basic *)sf->data;
11: PetscInt *rlengths, *ilengths, i, nRemoteRootRanks, nRemoteLeafRanks;
12: PetscMPIInt rank, niranks, *iranks, tag;
13: MPI_Comm comm;
14: MPI_Group group;
15: MPI_Request *rootreqs, *leafreqs;
17: MPI_Comm_group(PETSC_COMM_SELF, &group);
18: PetscSFSetUpRanks(sf, group);
19: MPI_Group_free(&group);
20: PetscObjectGetComm((PetscObject)sf, &comm);
21: PetscObjectGetNewTag((PetscObject)sf, &tag);
22: MPI_Comm_rank(comm, &rank);
23: /*
24: * Inform roots about how many leaves and from which ranks
25: */
26: PetscMalloc1(sf->nranks, &rlengths);
27: /* Determine number, sending ranks and length of incoming */
28: for (i = 0; i < sf->nranks; i++) { rlengths[i] = sf->roffset[i + 1] - sf->roffset[i]; /* Number of roots referenced by my leaves; for rank sf->ranks[i] */ }
29: nRemoteRootRanks = sf->nranks - sf->ndranks;
30: PetscCommBuildTwoSided(comm, 1, MPIU_INT, nRemoteRootRanks, sf->ranks + sf->ndranks, rlengths + sf->ndranks, &niranks, &iranks, (void **)&ilengths);
32: /* Sort iranks. See use of VecScatterGetRemoteOrdered_Private() in MatGetBrowsOfAoCols_MPIAIJ() on why.
33: We could sort ranks there at the price of allocating extra working arrays. Presumably, niranks is
34: small and the sorting is cheap.
35: */
36: PetscSortMPIIntWithIntArray(niranks, iranks, ilengths);
38: /* Partition into distinguished and non-distinguished incoming ranks */
39: bas->ndiranks = sf->ndranks;
40: bas->niranks = bas->ndiranks + niranks;
41: PetscMalloc2(bas->niranks, &bas->iranks, bas->niranks + 1, &bas->ioffset);
42: bas->ioffset[0] = 0;
43: for (i = 0; i < bas->ndiranks; i++) {
44: bas->iranks[i] = sf->ranks[i];
45: bas->ioffset[i + 1] = bas->ioffset[i] + rlengths[i];
46: }
48: for (; i < bas->niranks; i++) {
49: bas->iranks[i] = iranks[i - bas->ndiranks];
50: bas->ioffset[i + 1] = bas->ioffset[i] + ilengths[i - bas->ndiranks];
51: }
52: bas->itotal = bas->ioffset[i];
53: PetscFree(rlengths);
54: PetscFree(iranks);
55: PetscFree(ilengths);
57: /* Send leaf identities to roots */
58: nRemoteLeafRanks = bas->niranks - bas->ndiranks;
59: PetscMalloc1(bas->itotal, &bas->irootloc);
60: PetscMalloc2(nRemoteLeafRanks, &rootreqs, nRemoteRootRanks, &leafreqs);
61: for (i = bas->ndiranks; i < bas->niranks; i++) MPIU_Irecv(bas->irootloc + bas->ioffset[i], bas->ioffset[i + 1] - bas->ioffset[i], MPIU_INT, bas->iranks[i], tag, comm, &rootreqs[i - bas->ndiranks]);
62: for (i = 0; i < sf->nranks; i++) {
63: PetscInt npoints = sf->roffset[i + 1] - sf->roffset[i];
64: if (i < sf->ndranks) {
68: PetscArraycpy(bas->irootloc + bas->ioffset[0], sf->rremote + sf->roffset[i], npoints);
69: continue;
70: }
71: MPIU_Isend(sf->rremote + sf->roffset[i], npoints, MPIU_INT, sf->ranks[i], tag, comm, &leafreqs[i - sf->ndranks]);
72: }
73: MPI_Waitall(nRemoteLeafRanks, rootreqs, MPI_STATUSES_IGNORE);
74: MPI_Waitall(nRemoteRootRanks, leafreqs, MPI_STATUSES_IGNORE);
76: sf->nleafreqs = nRemoteRootRanks;
77: bas->nrootreqs = nRemoteLeafRanks;
78: sf->persistent = PETSC_TRUE;
80: /* Setup fields related to packing, such as rootbuflen[] */
81: PetscSFSetUpPackFields(sf);
82: PetscFree2(rootreqs, leafreqs);
83: return 0;
84: }
86: PETSC_INTERN PetscErrorCode PetscSFReset_Basic(PetscSF sf)
87: {
88: PetscSF_Basic *bas = (PetscSF_Basic *)sf->data;
89: PetscSFLink link = bas->avail, next;
92: PetscFree2(bas->iranks, bas->ioffset);
93: PetscFree(bas->irootloc);
95: #if defined(PETSC_HAVE_DEVICE)
96: for (PetscInt i = 0; i < 2; i++) PetscSFFree(sf, PETSC_MEMTYPE_DEVICE, bas->irootloc_d[i]);
97: #endif
99: #if defined(PETSC_HAVE_NVSHMEM)
100: PetscSFReset_Basic_NVSHMEM(sf);
101: #endif
103: for (; link; link = next) {
104: next = link->next;
105: PetscSFLinkDestroy(sf, link);
106: }
107: bas->avail = NULL;
108: PetscSFResetPackFields(sf);
109: return 0;
110: }
112: PETSC_INTERN PetscErrorCode PetscSFDestroy_Basic(PetscSF sf)
113: {
114: PetscSFReset_Basic(sf);
115: PetscFree(sf->data);
116: return 0;
117: }
119: #if defined(PETSC_USE_SINGLE_LIBRARY)
120: #include <petscmat.h>
122: PETSC_INTERN PetscErrorCode PetscSFView_Basic_PatternAndSizes(PetscSF sf, PetscViewer viewer)
123: {
124: PetscSF_Basic *bas = (PetscSF_Basic *)sf->data;
125: PetscInt i, nrootranks, ndrootranks;
126: const PetscInt *rootoffset;
127: PetscMPIInt rank, size;
128: const PetscMPIInt *rootranks;
129: MPI_Comm comm = PetscObjectComm((PetscObject)sf);
130: PetscScalar unitbytes;
131: Mat A;
133: MPI_Comm_size(comm, &size);
134: MPI_Comm_rank(comm, &rank);
135: /* PetscSFView is most useful for the SF used in VecScatterBegin/End in MatMult etc, where we do
136: PetscSFBcast, i.e., roots send data to leaves. We dump the communication pattern into a matrix
137: in senders' view point: how many bytes I will send to my neighbors.
139: Looking at a column of the matrix, one can also know how many bytes the rank will receive from others.
141: If PetscSFLink bas->inuse is available, we can use that to get tree vertex size. But that would give
142: different interpretations for the same SF for different data types. Since we most care about VecScatter,
143: we uniformly treat each vertex as a PetscScalar.
144: */
145: unitbytes = (PetscScalar)sizeof(PetscScalar);
147: PetscSFGetRootInfo_Basic(sf, &nrootranks, &ndrootranks, &rootranks, &rootoffset, NULL);
148: MatCreateAIJ(comm, 1, 1, size, size, 1, NULL, nrootranks - ndrootranks, NULL, &A);
149: MatSetOptionsPrefix(A, "__petsc_internal__"); /* To prevent the internal A from taking any command line options */
150: for (i = 0; i < nrootranks; i++) MatSetValue(A, (PetscInt)rank, bas->iranks[i], (rootoffset[i + 1] - rootoffset[i]) * unitbytes, INSERT_VALUES);
151: MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY);
152: MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);
153: MatView(A, viewer);
154: MatDestroy(&A);
155: return 0;
156: }
157: #endif
159: PETSC_INTERN PetscErrorCode PetscSFView_Basic(PetscSF sf, PetscViewer viewer)
160: {
161: PetscBool isascii;
163: PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii);
164: if (isascii && viewer->format != PETSC_VIEWER_ASCII_MATLAB) PetscViewerASCIIPrintf(viewer, " MultiSF sort=%s\n", sf->rankorder ? "rank-order" : "unordered");
165: #if defined(PETSC_USE_SINGLE_LIBRARY)
166: else {
167: PetscBool isdraw, isbinary;
168: PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw);
169: PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary);
170: if ((isascii && viewer->format == PETSC_VIEWER_ASCII_MATLAB) || isdraw || isbinary) PetscSFView_Basic_PatternAndSizes(sf, viewer);
171: }
172: #endif
173: return 0;
174: }
176: static PetscErrorCode PetscSFBcastBegin_Basic(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, const void *rootdata, PetscMemType leafmtype, void *leafdata, MPI_Op op)
177: {
178: PetscSFLink link = NULL;
180: /* Create a communication link, which provides buffers, MPI requests etc (if MPI is used) */
181: PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_BCAST, &link);
182: /* Pack rootdata to rootbuf for remote communication */
183: PetscSFLinkPackRootData(sf, link, PETSCSF_REMOTE, rootdata);
184: /* Start communcation, e.g., post MPI_Isend */
185: PetscSFLinkStartCommunication(sf, link, PETSCSF_../../../../../..2LEAF);
186: /* Do local scatter (i.e., self to self communication), which overlaps with the remote communication above */
187: PetscSFLinkScatterLocal(sf, link, PETSCSF_../../../../../..2LEAF, (void *)rootdata, leafdata, op);
188: return 0;
189: }
191: PETSC_INTERN PetscErrorCode PetscSFBcastEnd_Basic(PetscSF sf, MPI_Datatype unit, const void *rootdata, void *leafdata, MPI_Op op)
192: {
193: PetscSFLink link = NULL;
195: /* Retrieve the link used in XxxBegin() with root/leafdata as key */
196: PetscSFLinkGetInUse(sf, unit, rootdata, leafdata, PETSC_OWN_POINTER, &link);
197: /* Finish remote communication, e.g., post MPI_Waitall */
198: PetscSFLinkFinishCommunication(sf, link, PETSCSF_../../../../../..2LEAF);
199: /* Unpack data in leafbuf to leafdata for remote communication */
200: PetscSFLinkUnpackLeafData(sf, link, PETSCSF_REMOTE, leafdata, op);
201: /* Recycle the link */
202: PetscSFLinkReclaim(sf, &link);
203: return 0;
204: }
206: /* Shared by ReduceBegin and FetchAndOpBegin */
207: static inline PetscErrorCode PetscSFLeafToRootBegin_Basic(PetscSF sf, MPI_Datatype unit, PetscMemType leafmtype, const void *leafdata, PetscMemType rootmtype, void *rootdata, MPI_Op op, PetscSFOperation sfop, PetscSFLink *out)
208: {
209: PetscSFLink link = NULL;
211: PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, sfop, &link);
212: PetscSFLinkPackLeafData(sf, link, PETSCSF_REMOTE, leafdata);
213: PetscSFLinkStartCommunication(sf, link, PETSCSF_LEAF2../../../../../..);
214: *out = link;
215: return 0;
216: }
218: /* leaf -> root with reduction */
219: static PetscErrorCode PetscSFReduceBegin_Basic(PetscSF sf, MPI_Datatype unit, PetscMemType leafmtype, const void *leafdata, PetscMemType rootmtype, void *rootdata, MPI_Op op)
220: {
221: PetscSFLink link = NULL;
223: PetscSFLeafToRootBegin_Basic(sf, unit, leafmtype, leafdata, rootmtype, rootdata, op, PETSCSF_REDUCE, &link);
224: PetscSFLinkScatterLocal(sf, link, PETSCSF_LEAF2../../../../../.., rootdata, (void *)leafdata, op);
225: return 0;
226: }
228: PETSC_INTERN PetscErrorCode PetscSFReduceEnd_Basic(PetscSF sf, MPI_Datatype unit, const void *leafdata, void *rootdata, MPI_Op op)
229: {
230: PetscSFLink link = NULL;
232: PetscSFLinkGetInUse(sf, unit, rootdata, leafdata, PETSC_OWN_POINTER, &link);
233: PetscSFLinkFinishCommunication(sf, link, PETSCSF_LEAF2../../../../../..);
234: PetscSFLinkUnpackRootData(sf, link, PETSCSF_REMOTE, rootdata, op);
235: PetscSFLinkReclaim(sf, &link);
236: return 0;
237: }
239: PETSC_INTERN PetscErrorCode PetscSFFetchAndOpBegin_Basic(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, void *rootdata, PetscMemType leafmtype, const void *leafdata, void *leafupdate, MPI_Op op)
240: {
241: PetscSFLink link = NULL;
243: PetscSFLeafToRootBegin_Basic(sf, unit, leafmtype, leafdata, rootmtype, rootdata, op, PETSCSF_FETCH, &link);
244: PetscSFLinkFetchAndOpLocal(sf, link, rootdata, leafdata, leafupdate, op);
245: return 0;
246: }
248: static PetscErrorCode PetscSFFetchAndOpEnd_Basic(PetscSF sf, MPI_Datatype unit, void *rootdata, const void *leafdata, void *leafupdate, MPI_Op op)
249: {
250: PetscSFLink link = NULL;
252: PetscSFLinkGetInUse(sf, unit, rootdata, leafdata, PETSC_OWN_POINTER, &link);
253: /* This implementation could be changed to unpack as receives arrive, at the cost of non-determinism */
254: PetscSFLinkFinishCommunication(sf, link, PETSCSF_LEAF2../../../../../..);
255: /* Do fetch-and-op, the (remote) update results are in rootbuf */
256: PetscSFLinkFetchAndOpRemote(sf, link, rootdata, op);
257: /* Bcast rootbuf to leafupdate */
258: PetscSFLinkStartCommunication(sf, link, PETSCSF_../../../../../..2LEAF);
259: PetscSFLinkFinishCommunication(sf, link, PETSCSF_../../../../../..2LEAF);
260: /* Unpack and insert fetched data into leaves */
261: PetscSFLinkUnpackLeafData(sf, link, PETSCSF_REMOTE, leafupdate, MPI_REPLACE);
262: PetscSFLinkReclaim(sf, &link);
263: return 0;
264: }
266: PETSC_INTERN PetscErrorCode PetscSFGetLeafRanks_Basic(PetscSF sf, PetscInt *niranks, const PetscMPIInt **iranks, const PetscInt **ioffset, const PetscInt **irootloc)
267: {
268: PetscSF_Basic *bas = (PetscSF_Basic *)sf->data;
270: if (niranks) *niranks = bas->niranks;
271: if (iranks) *iranks = bas->iranks;
272: if (ioffset) *ioffset = bas->ioffset;
273: if (irootloc) *irootloc = bas->irootloc;
274: return 0;
275: }
277: /* An optimized PetscSFCreateEmbeddedRootSF. We aggresively make use of the established communication on sf.
278: We need one bcast on sf, and no communication anymore to build the embedded sf. Note that selected[]
279: was sorted before calling the routine.
280: */
281: PETSC_INTERN PetscErrorCode PetscSFCreateEmbeddedRootSF_Basic(PetscSF sf, PetscInt nselected, const PetscInt *selected, PetscSF *newsf)
282: {
283: PetscSF esf;
284: PetscInt esf_nranks, esf_ndranks, *esf_roffset, *esf_rmine, *esf_rremote;
285: PetscInt i, j, p, q, nroots, esf_nleaves, *new_ilocal, nranks, ndranks, niranks, ndiranks, minleaf, maxleaf, maxlocal;
286: char *rootdata, *leafdata, *leafmem; /* Only stores 0 or 1, so we can save memory with char */
287: PetscMPIInt *esf_ranks;
288: const PetscMPIInt *ranks, *iranks;
289: const PetscInt *roffset, *rmine, *rremote, *ioffset, *irootloc;
290: PetscBool connected;
291: PetscSFNode *new_iremote;
292: PetscSF_Basic *bas;
294: PetscSFCreate(PetscObjectComm((PetscObject)sf), &esf);
295: PetscSFSetFromOptions(esf);
296: PetscSFSetType(esf, PETSCSFBASIC); /* This optimized routine can only create a basic sf */
298: /* Find out which leaves are still connected to roots in the embedded sf by doing a Bcast */
299: PetscSFGetGraph(sf, &nroots, NULL, NULL, NULL);
300: PetscSFGetLeafRange(sf, &minleaf, &maxleaf);
301: maxlocal = maxleaf - minleaf + 1;
302: PetscCalloc2(nroots, &rootdata, maxlocal, &leafmem);
303: leafdata = leafmem - minleaf;
304: /* Tag selected roots */
305: for (i = 0; i < nselected; ++i) rootdata[selected[i]] = 1;
307: PetscSFBcastBegin(sf, MPI_CHAR, rootdata, leafdata, MPI_REPLACE);
308: PetscSFBcastEnd(sf, MPI_CHAR, rootdata, leafdata, MPI_REPLACE);
309: PetscSFGetLeafInfo_Basic(sf, &nranks, &ndranks, &ranks, &roffset, &rmine, &rremote); /* Get send info */
310: esf_nranks = esf_ndranks = esf_nleaves = 0;
311: for (i = 0; i < nranks; i++) {
312: connected = PETSC_FALSE; /* Is this process still connected to this remote root rank? */
313: for (j = roffset[i]; j < roffset[i + 1]; j++) {
314: if (leafdata[rmine[j]]) {
315: esf_nleaves++;
316: connected = PETSC_TRUE;
317: }
318: }
319: if (connected) {
320: esf_nranks++;
321: if (i < ndranks) esf_ndranks++;
322: }
323: }
325: /* Set graph of esf and also set up its outgoing communication (i.e., send info), which is usually done by PetscSFSetUpRanks */
326: PetscMalloc1(esf_nleaves, &new_ilocal);
327: PetscMalloc1(esf_nleaves, &new_iremote);
328: PetscMalloc4(esf_nranks, &esf_ranks, esf_nranks + 1, &esf_roffset, esf_nleaves, &esf_rmine, esf_nleaves, &esf_rremote);
329: p = 0; /* Counter for connected root ranks */
330: q = 0; /* Counter for connected leaves */
331: esf_roffset[0] = 0;
332: for (i = 0; i < nranks; i++) { /* Scan leaf data again to fill esf arrays */
333: connected = PETSC_FALSE;
334: for (j = roffset[i]; j < roffset[i + 1]; j++) {
335: if (leafdata[rmine[j]]) {
336: esf_rmine[q] = new_ilocal[q] = rmine[j];
337: esf_rremote[q] = rremote[j];
338: new_iremote[q].index = rremote[j];
339: new_iremote[q].rank = ranks[i];
340: connected = PETSC_TRUE;
341: q++;
342: }
343: }
344: if (connected) {
345: esf_ranks[p] = ranks[i];
346: esf_roffset[p + 1] = q;
347: p++;
348: }
349: }
351: /* SetGraph internally resets the SF, so we only set its fields after the call */
352: PetscSFSetGraph(esf, nroots, esf_nleaves, new_ilocal, PETSC_OWN_POINTER, new_iremote, PETSC_OWN_POINTER);
353: esf->nranks = esf_nranks;
354: esf->ndranks = esf_ndranks;
355: esf->ranks = esf_ranks;
356: esf->roffset = esf_roffset;
357: esf->rmine = esf_rmine;
358: esf->rremote = esf_rremote;
359: esf->nleafreqs = esf_nranks - esf_ndranks;
361: /* Set up the incoming communication (i.e., recv info) stored in esf->data, which is usually done by PetscSFSetUp_Basic */
362: bas = (PetscSF_Basic *)esf->data;
363: PetscSFGetRootInfo_Basic(sf, &niranks, &ndiranks, &iranks, &ioffset, &irootloc); /* Get recv info */
364: /* Embedded sf always has simpler communication than the original one. We might allocate longer arrays than needed here. But we
365: we do not care since these arrays are usually short. The benefit is we can fill these arrays by just parsing irootloc once.
366: */
367: PetscMalloc2(niranks, &bas->iranks, niranks + 1, &bas->ioffset);
368: PetscMalloc1(ioffset[niranks], &bas->irootloc);
369: bas->niranks = bas->ndiranks = bas->ioffset[0] = 0;
370: p = 0; /* Counter for connected leaf ranks */
371: q = 0; /* Counter for connected roots */
372: for (i = 0; i < niranks; i++) {
373: connected = PETSC_FALSE; /* Is the current process still connected to this remote leaf rank? */
374: for (j = ioffset[i]; j < ioffset[i + 1]; j++) {
375: if (rootdata[irootloc[j]]) {
376: bas->irootloc[q++] = irootloc[j];
377: connected = PETSC_TRUE;
378: }
379: }
380: if (connected) {
381: bas->niranks++;
382: if (i < ndiranks) bas->ndiranks++; /* Note that order of ranks (including distinguished ranks) is kept */
383: bas->iranks[p] = iranks[i];
384: bas->ioffset[p + 1] = q;
385: p++;
386: }
387: }
388: bas->itotal = q;
389: bas->nrootreqs = bas->niranks - bas->ndiranks;
390: esf->persistent = PETSC_TRUE;
391: /* Setup packing related fields */
392: PetscSFSetUpPackFields(esf);
394: /* Copy from PetscSFSetUp(), since this method wants to skip PetscSFSetUp(). */
395: #if defined(PETSC_HAVE_CUDA)
396: if (esf->backend == PETSCSF_BACKEND_CUDA) {
397: esf->ops->Malloc = PetscSFMalloc_CUDA;
398: esf->ops->Free = PetscSFFree_CUDA;
399: }
400: #endif
402: #if defined(PETSC_HAVE_HIP)
403: /* TODO: Needs debugging */
404: if (esf->backend == PETSCSF_BACKEND_HIP) {
405: esf->ops->Malloc = PetscSFMalloc_HIP;
406: esf->ops->Free = PetscSFFree_HIP;
407: }
408: #endif
410: #if defined(PETSC_HAVE_KOKKOS)
411: if (esf->backend == PETSCSF_BACKEND_KOKKOS) {
412: esf->ops->Malloc = PetscSFMalloc_Kokkos;
413: esf->ops->Free = PetscSFFree_Kokkos;
414: }
415: #endif
416: esf->setupcalled = PETSC_TRUE; /* We have done setup ourselves! */
417: PetscFree2(rootdata, leafmem);
418: *newsf = esf;
419: return 0;
420: }
422: PETSC_EXTERN PetscErrorCode PetscSFCreate_Basic(PetscSF sf)
423: {
424: PetscSF_Basic *dat;
426: sf->ops->SetUp = PetscSFSetUp_Basic;
427: sf->ops->Reset = PetscSFReset_Basic;
428: sf->ops->Destroy = PetscSFDestroy_Basic;
429: sf->ops->View = PetscSFView_Basic;
430: sf->ops->BcastBegin = PetscSFBcastBegin_Basic;
431: sf->ops->BcastEnd = PetscSFBcastEnd_Basic;
432: sf->ops->ReduceBegin = PetscSFReduceBegin_Basic;
433: sf->ops->ReduceEnd = PetscSFReduceEnd_Basic;
434: sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Basic;
435: sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Basic;
436: sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Basic;
437: sf->ops->CreateEmbeddedRootSF = PetscSFCreateEmbeddedRootSF_Basic;
439: PetscNew(&dat);
440: sf->data = (void *)dat;
441: return 0;
442: }