xref: /petsc/src/dm/impls/plex/plexpartition.c (revision 34c645fd3b0199e05bec2fcc32d3597bfeb7f4f2)
1 #include <petsc/private/dmpleximpl.h> /*I      "petscdmplex.h"   I*/
2 #include <petsc/private/partitionerimpl.h>
3 #include <petsc/private/hashseti.h>
4 
5 const char *const DMPlexCSRAlgorithms[] = {"mat", "graph", "overlap", "DMPlexCSRAlgorithm", "DM_PLEX_CSR_", NULL};
6 
7 static inline PetscInt DMPlex_GlobalID(PetscInt point)
8 {
9   return point >= 0 ? point : -(point + 1);
10 }
11 
12 static PetscErrorCode DMPlexCreatePartitionerGraph_Overlap(DM dm, PetscInt height, PetscInt *numVertices, PetscInt **offsets, PetscInt **adjacency, IS *globalNumbering)
13 {
14   DM              ovdm;
15   PetscSF         sfPoint;
16   IS              cellNumbering;
17   const PetscInt *cellNum;
18   PetscInt       *adj = NULL, *vOffsets = NULL, *vAdj = NULL;
19   PetscBool       useCone, useClosure;
20   PetscInt        dim, depth, overlap, cStart, cEnd, c, v;
21   PetscMPIInt     rank, size;
22 
23   PetscFunctionBegin;
24   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank));
25   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)dm), &size));
26   PetscCall(DMGetDimension(dm, &dim));
27   PetscCall(DMPlexGetDepth(dm, &depth));
28   if (dim != depth) {
29     /* We do not handle the uninterpolated case here */
30     PetscCall(DMPlexCreateNeighborCSR(dm, height, numVertices, offsets, adjacency));
31     /* DMPlexCreateNeighborCSR does not make a numbering */
32     if (globalNumbering) PetscCall(DMPlexCreateCellNumbering_Internal(dm, PETSC_TRUE, globalNumbering));
33     /* Different behavior for empty graphs */
34     if (!*numVertices) {
35       PetscCall(PetscMalloc1(1, offsets));
36       (*offsets)[0] = 0;
37     }
38     /* Broken in parallel */
39     if (rank) PetscCheck(!*numVertices, PETSC_COMM_SELF, PETSC_ERR_SUP, "Parallel partitioning of uninterpolated meshes not supported");
40     PetscFunctionReturn(PETSC_SUCCESS);
41   }
42   /* Always use FVM adjacency to create partitioner graph */
43   PetscCall(DMGetBasicAdjacency(dm, &useCone, &useClosure));
44   PetscCall(DMSetBasicAdjacency(dm, PETSC_TRUE, PETSC_FALSE));
45   /* Need overlap >= 1 */
46   PetscCall(DMPlexGetOverlap(dm, &overlap));
47   if (size && overlap < 1) {
48     PetscCall(DMPlexDistributeOverlap(dm, 1, NULL, &ovdm));
49   } else {
50     PetscCall(PetscObjectReference((PetscObject)dm));
51     ovdm = dm;
52   }
53   PetscCall(DMGetPointSF(ovdm, &sfPoint));
54   PetscCall(DMPlexGetHeightStratum(ovdm, height, &cStart, &cEnd));
55   PetscCall(DMPlexCreateNumbering_Plex(ovdm, cStart, cEnd, 0, NULL, sfPoint, &cellNumbering));
56   if (globalNumbering) {
57     PetscCall(PetscObjectReference((PetscObject)cellNumbering));
58     *globalNumbering = cellNumbering;
59   }
60   PetscCall(ISGetIndices(cellNumbering, &cellNum));
61   /* Determine sizes */
62   for (*numVertices = 0, c = cStart; c < cEnd; ++c) {
63     /* Skip non-owned cells in parallel (ParMetis expects no overlap) */
64     if (cellNum[c - cStart] < 0) continue;
65     (*numVertices)++;
66   }
67   PetscCall(PetscCalloc1(*numVertices + 1, &vOffsets));
68   for (c = cStart, v = 0; c < cEnd; ++c) {
69     PetscInt adjSize = PETSC_DETERMINE, a, vsize = 0;
70 
71     if (cellNum[c - cStart] < 0) continue;
72     PetscCall(DMPlexGetAdjacency(ovdm, c, &adjSize, &adj));
73     for (a = 0; a < adjSize; ++a) {
74       const PetscInt point = adj[a];
75       if (point != c && cStart <= point && point < cEnd) ++vsize;
76     }
77     vOffsets[v + 1] = vOffsets[v] + vsize;
78     ++v;
79   }
80   /* Determine adjacency */
81   PetscCall(PetscMalloc1(vOffsets[*numVertices], &vAdj));
82   for (c = cStart, v = 0; c < cEnd; ++c) {
83     PetscInt adjSize = PETSC_DETERMINE, a, off = vOffsets[v];
84 
85     /* Skip non-owned cells in parallel (ParMetis expects no overlap) */
86     if (cellNum[c - cStart] < 0) continue;
87     PetscCall(DMPlexGetAdjacency(ovdm, c, &adjSize, &adj));
88     for (a = 0; a < adjSize; ++a) {
89       const PetscInt point = adj[a];
90       if (point != c && cStart <= point && point < cEnd) vAdj[off++] = DMPlex_GlobalID(cellNum[point - cStart]);
91     }
92     PetscCheck(off == vOffsets[v + 1], PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Offsets %" PetscInt_FMT " should be %" PetscInt_FMT, off, vOffsets[v + 1]);
93     /* Sort adjacencies (not strictly necessary) */
94     PetscCall(PetscSortInt(off - vOffsets[v], &vAdj[vOffsets[v]]));
95     ++v;
96   }
97   PetscCall(PetscFree(adj));
98   PetscCall(ISRestoreIndices(cellNumbering, &cellNum));
99   PetscCall(ISDestroy(&cellNumbering));
100   PetscCall(DMSetBasicAdjacency(dm, useCone, useClosure));
101   PetscCall(DMDestroy(&ovdm));
102   if (offsets) {
103     *offsets = vOffsets;
104   } else PetscCall(PetscFree(vOffsets));
105   if (adjacency) {
106     *adjacency = vAdj;
107   } else PetscCall(PetscFree(vAdj));
108   PetscFunctionReturn(PETSC_SUCCESS);
109 }
110 
111 static PetscErrorCode DMPlexCreatePartitionerGraph_Native(DM dm, PetscInt height, PetscInt *numVertices, PetscInt **offsets, PetscInt **adjacency, IS *globalNumbering)
112 {
113   PetscInt        dim, depth, p, pStart, pEnd, a, adjSize, idx, size;
114   PetscInt       *adj = NULL, *vOffsets = NULL, *graph = NULL;
115   IS              cellNumbering;
116   const PetscInt *cellNum;
117   PetscBool       useCone, useClosure;
118   PetscSection    section;
119   PetscSegBuffer  adjBuffer;
120   PetscSF         sfPoint;
121   PetscInt       *adjCells = NULL, *remoteCells = NULL;
122   const PetscInt *local;
123   PetscInt        nroots, nleaves, l;
124   PetscMPIInt     rank;
125 
126   PetscFunctionBegin;
127   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank));
128   PetscCall(DMGetDimension(dm, &dim));
129   PetscCall(DMPlexGetDepth(dm, &depth));
130   if (dim != depth) {
131     /* We do not handle the uninterpolated case here */
132     PetscCall(DMPlexCreateNeighborCSR(dm, height, numVertices, offsets, adjacency));
133     /* DMPlexCreateNeighborCSR does not make a numbering */
134     if (globalNumbering) PetscCall(DMPlexCreateCellNumbering_Internal(dm, PETSC_TRUE, globalNumbering));
135     /* Different behavior for empty graphs */
136     if (!*numVertices) {
137       PetscCall(PetscMalloc1(1, offsets));
138       (*offsets)[0] = 0;
139     }
140     /* Broken in parallel */
141     if (rank) PetscCheck(!*numVertices, PETSC_COMM_SELF, PETSC_ERR_SUP, "Parallel partitioning of uninterpolated meshes not supported");
142     PetscFunctionReturn(PETSC_SUCCESS);
143   }
144   PetscCall(DMGetPointSF(dm, &sfPoint));
145   PetscCall(DMPlexGetHeightStratum(dm, height, &pStart, &pEnd));
146   /* Build adjacency graph via a section/segbuffer */
147   PetscCall(PetscSectionCreate(PetscObjectComm((PetscObject)dm), &section));
148   PetscCall(PetscSectionSetChart(section, pStart, pEnd));
149   PetscCall(PetscSegBufferCreate(sizeof(PetscInt), 1000, &adjBuffer));
150   /* Always use FVM adjacency to create partitioner graph */
151   PetscCall(DMGetBasicAdjacency(dm, &useCone, &useClosure));
152   PetscCall(DMSetBasicAdjacency(dm, PETSC_TRUE, PETSC_FALSE));
153   PetscCall(DMPlexCreateNumbering_Plex(dm, pStart, pEnd, 0, NULL, sfPoint, &cellNumbering));
154   if (globalNumbering) {
155     PetscCall(PetscObjectReference((PetscObject)cellNumbering));
156     *globalNumbering = cellNumbering;
157   }
158   PetscCall(ISGetIndices(cellNumbering, &cellNum));
159   /* For all boundary faces (including faces adjacent to a ghost cell), record the local cell in adjCells
160      Broadcast adjCells to remoteCells (to get cells from roots) and Reduce adjCells to remoteCells (to get cells from leaves)
161    */
162   PetscCall(PetscSFGetGraph(sfPoint, &nroots, &nleaves, &local, NULL));
163   if (nroots >= 0) {
164     PetscInt fStart, fEnd, f;
165 
166     PetscCall(PetscCalloc2(nroots, &adjCells, nroots, &remoteCells));
167     PetscCall(DMPlexGetHeightStratum(dm, height + 1, &fStart, &fEnd));
168     for (l = 0; l < nroots; ++l) adjCells[l] = -3;
169     for (f = fStart; f < fEnd; ++f) {
170       const PetscInt *support;
171       PetscInt        supportSize;
172 
173       PetscCall(DMPlexGetSupport(dm, f, &support));
174       PetscCall(DMPlexGetSupportSize(dm, f, &supportSize));
175       if (supportSize == 1) adjCells[f] = DMPlex_GlobalID(cellNum[support[0] - pStart]);
176       else if (supportSize == 2) {
177         PetscCall(PetscFindInt(support[0], nleaves, local, &p));
178         if (p >= 0) adjCells[f] = DMPlex_GlobalID(cellNum[support[1] - pStart]);
179         PetscCall(PetscFindInt(support[1], nleaves, local, &p));
180         if (p >= 0) adjCells[f] = DMPlex_GlobalID(cellNum[support[0] - pStart]);
181       }
182       /* Handle non-conforming meshes */
183       if (supportSize > 2) {
184         PetscInt        numChildren, i;
185         const PetscInt *children;
186 
187         PetscCall(DMPlexGetTreeChildren(dm, f, &numChildren, &children));
188         for (i = 0; i < numChildren; ++i) {
189           const PetscInt child = children[i];
190           if (fStart <= child && child < fEnd) {
191             PetscCall(DMPlexGetSupport(dm, child, &support));
192             PetscCall(DMPlexGetSupportSize(dm, child, &supportSize));
193             if (supportSize == 1) adjCells[child] = DMPlex_GlobalID(cellNum[support[0] - pStart]);
194             else if (supportSize == 2) {
195               PetscCall(PetscFindInt(support[0], nleaves, local, &p));
196               if (p >= 0) adjCells[child] = DMPlex_GlobalID(cellNum[support[1] - pStart]);
197               PetscCall(PetscFindInt(support[1], nleaves, local, &p));
198               if (p >= 0) adjCells[child] = DMPlex_GlobalID(cellNum[support[0] - pStart]);
199             }
200           }
201         }
202       }
203     }
204     for (l = 0; l < nroots; ++l) remoteCells[l] = -1;
205     PetscCall(PetscSFBcastBegin(dm->sf, MPIU_INT, adjCells, remoteCells, MPI_REPLACE));
206     PetscCall(PetscSFBcastEnd(dm->sf, MPIU_INT, adjCells, remoteCells, MPI_REPLACE));
207     PetscCall(PetscSFReduceBegin(dm->sf, MPIU_INT, adjCells, remoteCells, MPI_MAX));
208     PetscCall(PetscSFReduceEnd(dm->sf, MPIU_INT, adjCells, remoteCells, MPI_MAX));
209   }
210   /* Combine local and global adjacencies */
211   for (*numVertices = 0, p = pStart; p < pEnd; p++) {
212     /* Skip non-owned cells in parallel (ParMetis expects no overlap) */
213     if (nroots > 0) {
214       if (cellNum[p - pStart] < 0) continue;
215     }
216     /* Add remote cells */
217     if (remoteCells) {
218       const PetscInt  gp = DMPlex_GlobalID(cellNum[p - pStart]);
219       PetscInt        coneSize, numChildren, c, i;
220       const PetscInt *cone, *children;
221 
222       PetscCall(DMPlexGetCone(dm, p, &cone));
223       PetscCall(DMPlexGetConeSize(dm, p, &coneSize));
224       for (c = 0; c < coneSize; ++c) {
225         const PetscInt point = cone[c];
226         if (remoteCells[point] >= 0 && remoteCells[point] != gp) {
227           PetscInt *PETSC_RESTRICT pBuf;
228           PetscCall(PetscSectionAddDof(section, p, 1));
229           PetscCall(PetscSegBufferGetInts(adjBuffer, 1, &pBuf));
230           *pBuf = remoteCells[point];
231         }
232         /* Handle non-conforming meshes */
233         PetscCall(DMPlexGetTreeChildren(dm, point, &numChildren, &children));
234         for (i = 0; i < numChildren; ++i) {
235           const PetscInt child = children[i];
236           if (remoteCells[child] >= 0 && remoteCells[child] != gp) {
237             PetscInt *PETSC_RESTRICT pBuf;
238             PetscCall(PetscSectionAddDof(section, p, 1));
239             PetscCall(PetscSegBufferGetInts(adjBuffer, 1, &pBuf));
240             *pBuf = remoteCells[child];
241           }
242         }
243       }
244     }
245     /* Add local cells */
246     adjSize = PETSC_DETERMINE;
247     PetscCall(DMPlexGetAdjacency(dm, p, &adjSize, &adj));
248     for (a = 0; a < adjSize; ++a) {
249       const PetscInt point = adj[a];
250       if (point != p && pStart <= point && point < pEnd) {
251         PetscInt *PETSC_RESTRICT pBuf;
252         PetscCall(PetscSectionAddDof(section, p, 1));
253         PetscCall(PetscSegBufferGetInts(adjBuffer, 1, &pBuf));
254         *pBuf = DMPlex_GlobalID(cellNum[point - pStart]);
255       }
256     }
257     (*numVertices)++;
258   }
259   PetscCall(PetscFree(adj));
260   PetscCall(PetscFree2(adjCells, remoteCells));
261   PetscCall(DMSetBasicAdjacency(dm, useCone, useClosure));
262 
263   /* Derive CSR graph from section/segbuffer */
264   PetscCall(PetscSectionSetUp(section));
265   PetscCall(PetscSectionGetStorageSize(section, &size));
266   PetscCall(PetscMalloc1(*numVertices + 1, &vOffsets));
267   for (idx = 0, p = pStart; p < pEnd; p++) {
268     if (nroots > 0) {
269       if (cellNum[p - pStart] < 0) continue;
270     }
271     PetscCall(PetscSectionGetOffset(section, p, &vOffsets[idx++]));
272   }
273   vOffsets[*numVertices] = size;
274   PetscCall(PetscSegBufferExtractAlloc(adjBuffer, &graph));
275 
276   if (nroots >= 0) {
277     /* Filter out duplicate edges using section/segbuffer */
278     PetscCall(PetscSectionReset(section));
279     PetscCall(PetscSectionSetChart(section, 0, *numVertices));
280     for (p = 0; p < *numVertices; p++) {
281       PetscInt start = vOffsets[p], end = vOffsets[p + 1];
282       PetscInt numEdges = end - start, *PETSC_RESTRICT edges;
283       PetscCall(PetscSortRemoveDupsInt(&numEdges, &graph[start]));
284       PetscCall(PetscSectionSetDof(section, p, numEdges));
285       PetscCall(PetscSegBufferGetInts(adjBuffer, numEdges, &edges));
286       PetscCall(PetscArraycpy(edges, &graph[start], numEdges));
287     }
288     PetscCall(PetscFree(vOffsets));
289     PetscCall(PetscFree(graph));
290     /* Derive CSR graph from section/segbuffer */
291     PetscCall(PetscSectionSetUp(section));
292     PetscCall(PetscSectionGetStorageSize(section, &size));
293     PetscCall(PetscMalloc1(*numVertices + 1, &vOffsets));
294     for (idx = 0, p = 0; p < *numVertices; p++) PetscCall(PetscSectionGetOffset(section, p, &vOffsets[idx++]));
295     vOffsets[*numVertices] = size;
296     PetscCall(PetscSegBufferExtractAlloc(adjBuffer, &graph));
297   } else {
298     /* Sort adjacencies (not strictly necessary) */
299     for (p = 0; p < *numVertices; p++) {
300       PetscInt start = vOffsets[p], end = vOffsets[p + 1];
301       PetscCall(PetscSortInt(end - start, &graph[start]));
302     }
303   }
304 
305   if (offsets) *offsets = vOffsets;
306   if (adjacency) *adjacency = graph;
307 
308   /* Cleanup */
309   PetscCall(PetscSegBufferDestroy(&adjBuffer));
310   PetscCall(PetscSectionDestroy(&section));
311   PetscCall(ISRestoreIndices(cellNumbering, &cellNum));
312   PetscCall(ISDestroy(&cellNumbering));
313   PetscFunctionReturn(PETSC_SUCCESS);
314 }
315 
316 static PetscErrorCode DMPlexCreatePartitionerGraph_ViaMat(DM dm, PetscInt height, PetscInt *numVertices, PetscInt **offsets, PetscInt **adjacency, IS *globalNumbering)
317 {
318   Mat             conn, CSR;
319   IS              fis, cis, cis_own;
320   PetscSF         sfPoint;
321   const PetscInt *rows, *cols, *ii, *jj;
322   PetscInt       *idxs, *idxs2;
323   PetscInt        dim, depth, floc, cloc, i, M, N, c, lm, m, cStart, cEnd, fStart, fEnd;
324   PetscMPIInt     rank;
325   PetscBool       flg;
326 
327   PetscFunctionBegin;
328   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank));
329   PetscCall(DMGetDimension(dm, &dim));
330   PetscCall(DMPlexGetDepth(dm, &depth));
331   if (dim != depth) {
332     /* We do not handle the uninterpolated case here */
333     PetscCall(DMPlexCreateNeighborCSR(dm, height, numVertices, offsets, adjacency));
334     /* DMPlexCreateNeighborCSR does not make a numbering */
335     if (globalNumbering) PetscCall(DMPlexCreateCellNumbering_Internal(dm, PETSC_TRUE, globalNumbering));
336     /* Different behavior for empty graphs */
337     if (!*numVertices) {
338       PetscCall(PetscMalloc1(1, offsets));
339       (*offsets)[0] = 0;
340     }
341     /* Broken in parallel */
342     if (rank) PetscCheck(!*numVertices, PETSC_COMM_SELF, PETSC_ERR_SUP, "Parallel partitioning of uninterpolated meshes not supported");
343     PetscFunctionReturn(PETSC_SUCCESS);
344   }
345   /* Interpolated and parallel case */
346 
347   /* numbering */
348   PetscCall(DMGetPointSF(dm, &sfPoint));
349   PetscCall(DMPlexGetHeightStratum(dm, height, &cStart, &cEnd));
350   PetscCall(DMPlexGetHeightStratum(dm, height + 1, &fStart, &fEnd));
351   PetscCall(DMPlexCreateNumbering_Plex(dm, cStart, cEnd, 0, &N, sfPoint, &cis));
352   PetscCall(DMPlexCreateNumbering_Plex(dm, fStart, fEnd, 0, &M, sfPoint, &fis));
353   if (globalNumbering) PetscCall(ISDuplicate(cis, globalNumbering));
354 
355   /* get positive global ids and local sizes for facets and cells */
356   PetscCall(ISGetLocalSize(fis, &m));
357   PetscCall(ISGetIndices(fis, &rows));
358   PetscCall(PetscMalloc1(m, &idxs));
359   for (i = 0, floc = 0; i < m; i++) {
360     const PetscInt p = rows[i];
361 
362     if (p < 0) {
363       idxs[i] = -(p + 1);
364     } else {
365       idxs[i] = p;
366       floc += 1;
367     }
368   }
369   PetscCall(ISRestoreIndices(fis, &rows));
370   PetscCall(ISDestroy(&fis));
371   PetscCall(ISCreateGeneral(PETSC_COMM_SELF, m, idxs, PETSC_OWN_POINTER, &fis));
372 
373   PetscCall(ISGetLocalSize(cis, &m));
374   PetscCall(ISGetIndices(cis, &cols));
375   PetscCall(PetscMalloc1(m, &idxs));
376   PetscCall(PetscMalloc1(m, &idxs2));
377   for (i = 0, cloc = 0; i < m; i++) {
378     const PetscInt p = cols[i];
379 
380     if (p < 0) {
381       idxs[i] = -(p + 1);
382     } else {
383       idxs[i]       = p;
384       idxs2[cloc++] = p;
385     }
386   }
387   PetscCall(ISRestoreIndices(cis, &cols));
388   PetscCall(ISDestroy(&cis));
389   PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)dm), m, idxs, PETSC_OWN_POINTER, &cis));
390   PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)dm), cloc, idxs2, PETSC_OWN_POINTER, &cis_own));
391 
392   /* Create matrix to hold F-C connectivity (MatMatTranspose Mult not supported for MPIAIJ) */
393   PetscCall(MatCreate(PetscObjectComm((PetscObject)dm), &conn));
394   PetscCall(MatSetSizes(conn, floc, cloc, M, N));
395   PetscCall(MatSetType(conn, MATMPIAIJ));
396   PetscCall(DMPlexGetMaxSizes(dm, NULL, &lm));
397   PetscCall(MPIU_Allreduce(&lm, &m, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)dm)));
398   PetscCall(MatMPIAIJSetPreallocation(conn, m, NULL, m, NULL));
399 
400   /* Assemble matrix */
401   PetscCall(ISGetIndices(fis, &rows));
402   PetscCall(ISGetIndices(cis, &cols));
403   for (c = cStart; c < cEnd; c++) {
404     const PetscInt *cone;
405     PetscInt        coneSize, row, col, f;
406 
407     col = cols[c - cStart];
408     PetscCall(DMPlexGetCone(dm, c, &cone));
409     PetscCall(DMPlexGetConeSize(dm, c, &coneSize));
410     for (f = 0; f < coneSize; f++) {
411       const PetscScalar v = 1.0;
412       const PetscInt   *children;
413       PetscInt          numChildren, ch;
414 
415       row = rows[cone[f] - fStart];
416       PetscCall(MatSetValues(conn, 1, &row, 1, &col, &v, INSERT_VALUES));
417 
418       /* non-conforming meshes */
419       PetscCall(DMPlexGetTreeChildren(dm, cone[f], &numChildren, &children));
420       for (ch = 0; ch < numChildren; ch++) {
421         const PetscInt child = children[ch];
422 
423         if (child < fStart || child >= fEnd) continue;
424         row = rows[child - fStart];
425         PetscCall(MatSetValues(conn, 1, &row, 1, &col, &v, INSERT_VALUES));
426       }
427     }
428   }
429   PetscCall(ISRestoreIndices(fis, &rows));
430   PetscCall(ISRestoreIndices(cis, &cols));
431   PetscCall(ISDestroy(&fis));
432   PetscCall(ISDestroy(&cis));
433   PetscCall(MatAssemblyBegin(conn, MAT_FINAL_ASSEMBLY));
434   PetscCall(MatAssemblyEnd(conn, MAT_FINAL_ASSEMBLY));
435 
436   /* Get parallel CSR by doing conn^T * conn */
437   PetscCall(MatTransposeMatMult(conn, conn, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &CSR));
438   PetscCall(MatDestroy(&conn));
439 
440   /* extract local part of the CSR */
441   PetscCall(MatMPIAIJGetLocalMat(CSR, MAT_INITIAL_MATRIX, &conn));
442   PetscCall(MatDestroy(&CSR));
443   PetscCall(MatGetRowIJ(conn, 0, PETSC_FALSE, PETSC_FALSE, &m, &ii, &jj, &flg));
444   PetscCheck(flg, PETSC_COMM_SELF, PETSC_ERR_PLIB, "No IJ format");
445 
446   /* get back requested output */
447   if (numVertices) *numVertices = m;
448   if (offsets) {
449     PetscCall(PetscCalloc1(m + 1, &idxs));
450     for (i = 1; i < m + 1; i++) idxs[i] = ii[i] - i; /* ParMetis does not like self-connectivity */
451     *offsets = idxs;
452   }
453   if (adjacency) {
454     PetscCall(PetscMalloc1(ii[m] - m, &idxs));
455     PetscCall(ISGetIndices(cis_own, &rows));
456     for (i = 0, c = 0; i < m; i++) {
457       PetscInt j, g = rows[i];
458 
459       for (j = ii[i]; j < ii[i + 1]; j++) {
460         if (jj[j] == g) continue; /* again, self-connectivity */
461         idxs[c++] = jj[j];
462       }
463     }
464     PetscCheck(c == ii[m] - m, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Unexpected %" PetscInt_FMT " != %" PetscInt_FMT, c, ii[m] - m);
465     PetscCall(ISRestoreIndices(cis_own, &rows));
466     *adjacency = idxs;
467   }
468 
469   /* cleanup */
470   PetscCall(ISDestroy(&cis_own));
471   PetscCall(MatRestoreRowIJ(conn, 0, PETSC_FALSE, PETSC_FALSE, &m, &ii, &jj, &flg));
472   PetscCheck(flg, PETSC_COMM_SELF, PETSC_ERR_PLIB, "No IJ format");
473   PetscCall(MatDestroy(&conn));
474   PetscFunctionReturn(PETSC_SUCCESS);
475 }
476 
477 /*@C
478   DMPlexCreatePartitionerGraph - Create a CSR graph of point connections for the partitioner
479 
480   Collective
481 
482   Input Parameters:
483 + dm     - The mesh `DM`
484 - height - Height of the strata from which to construct the graph
485 
486   Output Parameters:
487 + numVertices     - Number of vertices in the graph
488 . offsets         - Point offsets in the graph
489 . adjacency       - Point connectivity in the graph
490 - globalNumbering - A map from the local cell numbering to the global numbering used in "adjacency".  Negative indicates that the cell is a duplicate from another process.
491 
492   Options Database Key:
493 . -dm_plex_csr_alg <mat,graph,overlap> - Choose the algorithm for computing the CSR graph
494 
495   Level: developer
496 
497   Note:
498   The user can control the definition of adjacency for the mesh using `DMSetAdjacency()`. They should choose the combination appropriate for the function
499   representation on the mesh. If requested, globalNumbering needs to be destroyed by the caller; offsets and adjacency need to be freed with PetscFree().
500 
501 .seealso: [](ch_unstructured), `DM`, `DMPLEX`, `PetscPartitionerGetType()`, `PetscPartitionerCreate()`, `DMSetAdjacency()`
502 @*/
503 PetscErrorCode DMPlexCreatePartitionerGraph(DM dm, PetscInt height, PetscInt *numVertices, PetscInt **offsets, PetscInt **adjacency, IS *globalNumbering)
504 {
505   DMPlexCSRAlgorithm alg = DM_PLEX_CSR_GRAPH;
506 
507   PetscFunctionBegin;
508   PetscCall(PetscOptionsGetEnum(((PetscObject)dm)->options, ((PetscObject)dm)->prefix, "-dm_plex_csr_alg", DMPlexCSRAlgorithms, (PetscEnum *)&alg, NULL));
509   switch (alg) {
510   case DM_PLEX_CSR_MAT:
511     PetscCall(DMPlexCreatePartitionerGraph_ViaMat(dm, height, numVertices, offsets, adjacency, globalNumbering));
512     break;
513   case DM_PLEX_CSR_GRAPH:
514     PetscCall(DMPlexCreatePartitionerGraph_Native(dm, height, numVertices, offsets, adjacency, globalNumbering));
515     break;
516   case DM_PLEX_CSR_OVERLAP:
517     PetscCall(DMPlexCreatePartitionerGraph_Overlap(dm, height, numVertices, offsets, adjacency, globalNumbering));
518     break;
519   }
520   PetscFunctionReturn(PETSC_SUCCESS);
521 }
522 
523 /*@C
524   DMPlexCreateNeighborCSR - Create a mesh graph (cell-cell adjacency) in parallel CSR format.
525 
526   Collective
527 
528   Input Parameters:
529 + dm         - The `DMPLEX`
530 - cellHeight - The height of mesh points to treat as cells (default should be 0)
531 
532   Output Parameters:
533 + numVertices - The number of local vertices in the graph, or cells in the mesh.
534 . offsets     - The offset to the adjacency list for each cell
535 - adjacency   - The adjacency list for all cells
536 
537   Level: advanced
538 
539   Note:
540   This is suitable for input to a mesh partitioner like ParMetis.
541 
542 .seealso: [](ch_unstructured), `DM`, `DMPLEX`, `DMPlexCreate()`
543 @*/
544 PetscErrorCode DMPlexCreateNeighborCSR(DM dm, PetscInt cellHeight, PetscInt *numVertices, PetscInt **offsets, PetscInt **adjacency)
545 {
546   const PetscInt maxFaceCases = 30;
547   PetscInt       numFaceCases = 0;
548   PetscInt       numFaceVertices[30]; /* maxFaceCases, C89 sucks sucks sucks */
549   PetscInt      *off, *adj;
550   PetscInt      *neighborCells = NULL;
551   PetscInt       dim, cellDim, depth = 0, faceDepth, cStart, cEnd, c, numCells, cell;
552 
553   PetscFunctionBegin;
554   /* For parallel partitioning, I think you have to communicate supports */
555   PetscCall(DMGetDimension(dm, &dim));
556   cellDim = dim - cellHeight;
557   PetscCall(DMPlexGetDepth(dm, &depth));
558   PetscCall(DMPlexGetHeightStratum(dm, cellHeight, &cStart, &cEnd));
559   if (cEnd - cStart == 0) {
560     if (numVertices) *numVertices = 0;
561     if (offsets) *offsets = NULL;
562     if (adjacency) *adjacency = NULL;
563     PetscFunctionReturn(PETSC_SUCCESS);
564   }
565   numCells  = cEnd - cStart;
566   faceDepth = depth - cellHeight;
567   if (dim == depth) {
568     PetscInt f, fStart, fEnd;
569 
570     PetscCall(PetscCalloc1(numCells + 1, &off));
571     /* Count neighboring cells */
572     PetscCall(DMPlexGetHeightStratum(dm, cellHeight + 1, &fStart, &fEnd));
573     for (f = fStart; f < fEnd; ++f) {
574       const PetscInt *support;
575       PetscInt        supportSize;
576       PetscCall(DMPlexGetSupportSize(dm, f, &supportSize));
577       PetscCall(DMPlexGetSupport(dm, f, &support));
578       if (supportSize == 2) {
579         PetscInt numChildren;
580 
581         PetscCall(DMPlexGetTreeChildren(dm, f, &numChildren, NULL));
582         if (!numChildren) {
583           ++off[support[0] - cStart + 1];
584           ++off[support[1] - cStart + 1];
585         }
586       }
587     }
588     /* Prefix sum */
589     for (c = 1; c <= numCells; ++c) off[c] += off[c - 1];
590     if (adjacency) {
591       PetscInt *tmp;
592 
593       PetscCall(PetscMalloc1(off[numCells], &adj));
594       PetscCall(PetscMalloc1(numCells + 1, &tmp));
595       PetscCall(PetscArraycpy(tmp, off, numCells + 1));
596       /* Get neighboring cells */
597       for (f = fStart; f < fEnd; ++f) {
598         const PetscInt *support;
599         PetscInt        supportSize;
600         PetscCall(DMPlexGetSupportSize(dm, f, &supportSize));
601         PetscCall(DMPlexGetSupport(dm, f, &support));
602         if (supportSize == 2) {
603           PetscInt numChildren;
604 
605           PetscCall(DMPlexGetTreeChildren(dm, f, &numChildren, NULL));
606           if (!numChildren) {
607             adj[tmp[support[0] - cStart]++] = support[1];
608             adj[tmp[support[1] - cStart]++] = support[0];
609           }
610         }
611       }
612       for (c = 0; c < cEnd - cStart; ++c) PetscAssert(tmp[c] == off[c + 1], PETSC_COMM_SELF, PETSC_ERR_PLIB, "Offset %" PetscInt_FMT " != %" PetscInt_FMT " for cell %" PetscInt_FMT, tmp[c], off[c], c + cStart);
613       PetscCall(PetscFree(tmp));
614     }
615     if (numVertices) *numVertices = numCells;
616     if (offsets) *offsets = off;
617     if (adjacency) *adjacency = adj;
618     PetscFunctionReturn(PETSC_SUCCESS);
619   }
620   /* Setup face recognition */
621   if (faceDepth == 1) {
622     PetscInt cornersSeen[30] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; /* Could use PetscBT */
623 
624     for (c = cStart; c < cEnd; ++c) {
625       PetscInt corners;
626 
627       PetscCall(DMPlexGetConeSize(dm, c, &corners));
628       if (!cornersSeen[corners]) {
629         PetscInt nFV;
630 
631         PetscCheck(numFaceCases < maxFaceCases, PetscObjectComm((PetscObject)dm), PETSC_ERR_PLIB, "Exceeded maximum number of face recognition cases");
632         cornersSeen[corners] = 1;
633 
634         PetscCall(DMPlexGetNumFaceVertices(dm, cellDim, corners, &nFV));
635 
636         numFaceVertices[numFaceCases++] = nFV;
637       }
638     }
639   }
640   PetscCall(PetscCalloc1(numCells + 1, &off));
641   /* Count neighboring cells */
642   for (cell = cStart; cell < cEnd; ++cell) {
643     PetscInt numNeighbors = PETSC_DETERMINE, n;
644 
645     PetscCall(DMPlexGetAdjacency_Internal(dm, cell, PETSC_TRUE, PETSC_FALSE, PETSC_FALSE, &numNeighbors, &neighborCells));
646     /* Get meet with each cell, and check with recognizer (could optimize to check each pair only once) */
647     for (n = 0; n < numNeighbors; ++n) {
648       PetscInt        cellPair[2];
649       PetscBool       found    = faceDepth > 1 ? PETSC_TRUE : PETSC_FALSE;
650       PetscInt        meetSize = 0;
651       const PetscInt *meet     = NULL;
652 
653       cellPair[0] = cell;
654       cellPair[1] = neighborCells[n];
655       if (cellPair[0] == cellPair[1]) continue;
656       if (!found) {
657         PetscCall(DMPlexGetMeet(dm, 2, cellPair, &meetSize, &meet));
658         if (meetSize) {
659           PetscInt f;
660 
661           for (f = 0; f < numFaceCases; ++f) {
662             if (numFaceVertices[f] == meetSize) {
663               found = PETSC_TRUE;
664               break;
665             }
666           }
667         }
668         PetscCall(DMPlexRestoreMeet(dm, 2, cellPair, &meetSize, &meet));
669       }
670       if (found) ++off[cell - cStart + 1];
671     }
672   }
673   /* Prefix sum */
674   for (cell = 1; cell <= numCells; ++cell) off[cell] += off[cell - 1];
675 
676   if (adjacency) {
677     PetscCall(PetscMalloc1(off[numCells], &adj));
678     /* Get neighboring cells */
679     for (cell = cStart; cell < cEnd; ++cell) {
680       PetscInt numNeighbors = PETSC_DETERMINE, n;
681       PetscInt cellOffset   = 0;
682 
683       PetscCall(DMPlexGetAdjacency_Internal(dm, cell, PETSC_TRUE, PETSC_FALSE, PETSC_FALSE, &numNeighbors, &neighborCells));
684       /* Get meet with each cell, and check with recognizer (could optimize to check each pair only once) */
685       for (n = 0; n < numNeighbors; ++n) {
686         PetscInt        cellPair[2];
687         PetscBool       found    = faceDepth > 1 ? PETSC_TRUE : PETSC_FALSE;
688         PetscInt        meetSize = 0;
689         const PetscInt *meet     = NULL;
690 
691         cellPair[0] = cell;
692         cellPair[1] = neighborCells[n];
693         if (cellPair[0] == cellPair[1]) continue;
694         if (!found) {
695           PetscCall(DMPlexGetMeet(dm, 2, cellPair, &meetSize, &meet));
696           if (meetSize) {
697             PetscInt f;
698 
699             for (f = 0; f < numFaceCases; ++f) {
700               if (numFaceVertices[f] == meetSize) {
701                 found = PETSC_TRUE;
702                 break;
703               }
704             }
705           }
706           PetscCall(DMPlexRestoreMeet(dm, 2, cellPair, &meetSize, &meet));
707         }
708         if (found) {
709           adj[off[cell - cStart] + cellOffset] = neighborCells[n];
710           ++cellOffset;
711         }
712       }
713     }
714   }
715   PetscCall(PetscFree(neighborCells));
716   if (numVertices) *numVertices = numCells;
717   if (offsets) *offsets = off;
718   if (adjacency) *adjacency = adj;
719   PetscFunctionReturn(PETSC_SUCCESS);
720 }
721 
722 /*@
723   PetscPartitionerDMPlexPartition - Create a non-overlapping partition of the cells in the mesh
724 
725   Collective
726 
727   Input Parameters:
728 + part          - The `PetscPartitioner`
729 . targetSection - The `PetscSection` describing the absolute weight of each partition (can be `NULL`)
730 - dm            - The mesh `DM`
731 
732   Output Parameters:
733 + partSection - The `PetscSection` giving the division of points by partition
734 - partition   - The list of points by partition
735 
736   Level: developer
737 
738   Note:
739   If the `DM` has a local section associated, each point to be partitioned will be weighted by the total number of dofs identified
740   by the section in the transitive closure of the point.
741 
742 .seealso: [](ch_unstructured), `DM`, `DMPLEX`, `PetscPartitioner`, `PetscSection`, `DMPlexDistribute()`, `PetscPartitionerCreate()`, `PetscSectionCreate()`,
743          `PetscSectionSetChart()`, `PetscPartitionerPartition()`
744 @*/
745 PetscErrorCode PetscPartitionerDMPlexPartition(PetscPartitioner part, DM dm, PetscSection targetSection, PetscSection partSection, IS *partition)
746 {
747   PetscMPIInt  size;
748   PetscBool    isplex;
749   PetscSection vertSection = NULL, edgeSection = NULL;
750 
751   PetscFunctionBegin;
752   PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1);
753   PetscValidHeaderSpecific(dm, DM_CLASSID, 2);
754   if (targetSection) PetscValidHeaderSpecific(targetSection, PETSC_SECTION_CLASSID, 3);
755   PetscValidHeaderSpecific(partSection, PETSC_SECTION_CLASSID, 4);
756   PetscAssertPointer(partition, 5);
757   PetscCall(PetscObjectTypeCompare((PetscObject)dm, DMPLEX, &isplex));
758   PetscCheck(isplex, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Not for type %s", ((PetscObject)dm)->type_name);
759   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)part), &size));
760   if (size == 1) {
761     PetscInt *points;
762     PetscInt  cStart, cEnd, c;
763 
764     PetscCall(DMPlexGetHeightStratum(dm, part->height, &cStart, &cEnd));
765     PetscCall(PetscSectionReset(partSection));
766     PetscCall(PetscSectionSetChart(partSection, 0, size));
767     PetscCall(PetscSectionSetDof(partSection, 0, cEnd - cStart));
768     PetscCall(PetscSectionSetUp(partSection));
769     PetscCall(PetscMalloc1(cEnd - cStart, &points));
770     for (c = cStart; c < cEnd; ++c) points[c] = c;
771     PetscCall(ISCreateGeneral(PetscObjectComm((PetscObject)part), cEnd - cStart, points, PETSC_OWN_POINTER, partition));
772     PetscFunctionReturn(PETSC_SUCCESS);
773   }
774   if (part->height == 0) {
775     PetscInt  numVertices = 0;
776     PetscInt *start       = NULL;
777     PetscInt *adjacency   = NULL;
778     IS        globalNumbering;
779 
780     if (!part->noGraph || part->viewGraph) {
781       PetscCall(DMPlexCreatePartitionerGraph(dm, part->height, &numVertices, &start, &adjacency, &globalNumbering));
782     } else { /* only compute the number of owned local vertices */
783       const PetscInt *idxs;
784       PetscInt        p, pStart, pEnd;
785 
786       PetscCall(DMPlexGetHeightStratum(dm, part->height, &pStart, &pEnd));
787       PetscCall(DMPlexCreateNumbering_Plex(dm, pStart, pEnd, 0, NULL, dm->sf, &globalNumbering));
788       PetscCall(ISGetIndices(globalNumbering, &idxs));
789       for (p = 0; p < pEnd - pStart; p++) numVertices += idxs[p] < 0 ? 0 : 1;
790       PetscCall(ISRestoreIndices(globalNumbering, &idxs));
791     }
792     if (part->usevwgt) {
793       PetscSection    section = dm->localSection, clSection = NULL;
794       IS              clPoints = NULL;
795       const PetscInt *gid, *clIdx;
796       PetscInt        v, p, pStart, pEnd;
797 
798       /* dm->localSection encodes degrees of freedom per point, not per cell. We need to get the closure index to properly specify cell weights (aka dofs) */
799       /* We do this only if the local section has been set */
800       if (section) {
801         PetscCall(PetscSectionGetClosureIndex(section, (PetscObject)dm, &clSection, NULL));
802         if (!clSection) PetscCall(DMPlexCreateClosureIndex(dm, NULL));
803         PetscCall(PetscSectionGetClosureIndex(section, (PetscObject)dm, &clSection, &clPoints));
804         PetscCall(ISGetIndices(clPoints, &clIdx));
805       }
806       PetscCall(DMPlexGetHeightStratum(dm, part->height, &pStart, &pEnd));
807       PetscCall(PetscSectionCreate(PETSC_COMM_SELF, &vertSection));
808       PetscCall(PetscSectionSetChart(vertSection, 0, numVertices));
809       if (globalNumbering) {
810         PetscCall(ISGetIndices(globalNumbering, &gid));
811       } else gid = NULL;
812       for (p = pStart, v = 0; p < pEnd; ++p) {
813         PetscInt dof = 1;
814 
815         /* skip cells in the overlap */
816         if (gid && gid[p - pStart] < 0) continue;
817 
818         if (section) {
819           PetscInt cl, clSize, clOff;
820 
821           dof = 0;
822           PetscCall(PetscSectionGetDof(clSection, p, &clSize));
823           PetscCall(PetscSectionGetOffset(clSection, p, &clOff));
824           for (cl = 0; cl < clSize; cl += 2) {
825             PetscInt clDof, clPoint = clIdx[clOff + cl]; /* odd indices are reserved for orientations */
826 
827             PetscCall(PetscSectionGetDof(section, clPoint, &clDof));
828             dof += clDof;
829           }
830         }
831         PetscCheck(dof, PETSC_COMM_SELF, PETSC_ERR_SUP, "Number of dofs for point %" PetscInt_FMT " in the local section should be positive", p);
832         PetscCall(PetscSectionSetDof(vertSection, v, dof));
833         v++;
834       }
835       if (globalNumbering) PetscCall(ISRestoreIndices(globalNumbering, &gid));
836       if (clPoints) PetscCall(ISRestoreIndices(clPoints, &clIdx));
837       PetscCall(PetscSectionSetUp(vertSection));
838     }
839     if (part->useewgt) {
840       const PetscInt numEdges = start[numVertices];
841 
842       PetscCall(PetscSectionCreate(PETSC_COMM_SELF, &edgeSection));
843       PetscCall(PetscSectionSetChart(edgeSection, 0, numEdges));
844       for (PetscInt e = 0; e < start[numVertices]; ++e) PetscCall(PetscSectionSetDof(edgeSection, e, 1));
845       for (PetscInt v = 0; v < numVertices; ++v) {
846         DMPolytopeType ct;
847 
848         // Assume v is the cell number
849         PetscCall(DMPlexGetCellType(dm, v, &ct));
850         if (ct != DM_POLYTOPE_POINT_PRISM_TENSOR && ct != DM_POLYTOPE_SEG_PRISM_TENSOR && ct != DM_POLYTOPE_TRI_PRISM_TENSOR && ct != DM_POLYTOPE_QUAD_PRISM_TENSOR) continue;
851 
852         for (PetscInt e = start[v]; e < start[v + 1]; ++e) PetscCall(PetscSectionSetDof(edgeSection, e, 3));
853       }
854       PetscCall(PetscSectionSetUp(edgeSection));
855     }
856     PetscCall(PetscPartitionerPartition(part, size, numVertices, start, adjacency, vertSection, edgeSection, targetSection, partSection, partition));
857     PetscCall(PetscFree(start));
858     PetscCall(PetscFree(adjacency));
859     if (globalNumbering) { /* partition is wrt global unique numbering: change this to be wrt local numbering */
860       const PetscInt *globalNum;
861       const PetscInt *partIdx;
862       PetscInt       *map, cStart, cEnd;
863       PetscInt       *adjusted, i, localSize, offset;
864       IS              newPartition;
865 
866       PetscCall(ISGetLocalSize(*partition, &localSize));
867       PetscCall(PetscMalloc1(localSize, &adjusted));
868       PetscCall(ISGetIndices(globalNumbering, &globalNum));
869       PetscCall(ISGetIndices(*partition, &partIdx));
870       PetscCall(PetscMalloc1(localSize, &map));
871       PetscCall(DMPlexGetHeightStratum(dm, part->height, &cStart, &cEnd));
872       for (i = cStart, offset = 0; i < cEnd; i++) {
873         if (globalNum[i - cStart] >= 0) map[offset++] = i;
874       }
875       for (i = 0; i < localSize; i++) adjusted[i] = map[partIdx[i]];
876       PetscCall(PetscFree(map));
877       PetscCall(ISRestoreIndices(*partition, &partIdx));
878       PetscCall(ISRestoreIndices(globalNumbering, &globalNum));
879       PetscCall(ISCreateGeneral(PETSC_COMM_SELF, localSize, adjusted, PETSC_OWN_POINTER, &newPartition));
880       PetscCall(ISDestroy(&globalNumbering));
881       PetscCall(ISDestroy(partition));
882       *partition = newPartition;
883     }
884   } else SETERRQ(PetscObjectComm((PetscObject)part), PETSC_ERR_ARG_OUTOFRANGE, "Invalid height %" PetscInt_FMT " for points to partition", part->height);
885   PetscCall(PetscSectionDestroy(&vertSection));
886   PetscCall(PetscSectionDestroy(&edgeSection));
887   PetscFunctionReturn(PETSC_SUCCESS);
888 }
889 
890 /*@
891   DMPlexGetPartitioner - Get the mesh partitioner
892 
893   Not Collective
894 
895   Input Parameter:
896 . dm - The `DM`
897 
898   Output Parameter:
899 . part - The `PetscPartitioner`
900 
901   Level: developer
902 
903   Note:
904   This gets a borrowed reference, so the user should not destroy this `PetscPartitioner`.
905 
906 .seealso: [](ch_unstructured), `DM`, `DMPLEX`, `PetscPartitioner`, `PetscSection`, `DMPlexDistribute()`, `DMPlexSetPartitioner()`, `PetscPartitionerDMPlexPartition()`, `PetscPartitionerCreate()`
907 @*/
908 PetscErrorCode DMPlexGetPartitioner(DM dm, PetscPartitioner *part)
909 {
910   DM_Plex *mesh = (DM_Plex *)dm->data;
911 
912   PetscFunctionBegin;
913   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
914   PetscAssertPointer(part, 2);
915   *part = mesh->partitioner;
916   PetscFunctionReturn(PETSC_SUCCESS);
917 }
918 
919 /*@
920   DMPlexSetPartitioner - Set the mesh partitioner
921 
922   logically Collective
923 
924   Input Parameters:
925 + dm   - The `DM`
926 - part - The partitioner
927 
928   Level: developer
929 
930   Note:
931   Any existing `PetscPartitioner` will be destroyed.
932 
933 .seealso: [](ch_unstructured), `DM`, `DMPLEX`, `PetscPartitioner`,`DMPlexDistribute()`, `DMPlexGetPartitioner()`, `PetscPartitionerCreate()`
934 @*/
935 PetscErrorCode DMPlexSetPartitioner(DM dm, PetscPartitioner part)
936 {
937   DM_Plex *mesh = (DM_Plex *)dm->data;
938 
939   PetscFunctionBegin;
940   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
941   PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 2);
942   PetscCall(PetscObjectReference((PetscObject)part));
943   PetscCall(PetscPartitionerDestroy(&mesh->partitioner));
944   mesh->partitioner = part;
945   PetscFunctionReturn(PETSC_SUCCESS);
946 }
947 
948 static PetscErrorCode DMPlexAddClosure_Private(DM dm, PetscHSetI ht, PetscInt point)
949 {
950   const PetscInt *cone;
951   PetscInt        coneSize, c;
952   PetscBool       missing;
953 
954   PetscFunctionBeginHot;
955   PetscCall(PetscHSetIQueryAdd(ht, point, &missing));
956   if (missing) {
957     PetscCall(DMPlexGetCone(dm, point, &cone));
958     PetscCall(DMPlexGetConeSize(dm, point, &coneSize));
959     for (c = 0; c < coneSize; c++) PetscCall(DMPlexAddClosure_Private(dm, ht, cone[c]));
960   }
961   PetscFunctionReturn(PETSC_SUCCESS);
962 }
963 
964 PETSC_UNUSED static PetscErrorCode DMPlexAddClosure_Tree(DM dm, PetscHSetI ht, PetscInt point, PetscBool up, PetscBool down)
965 {
966   PetscFunctionBegin;
967   if (up) {
968     PetscInt parent;
969 
970     PetscCall(DMPlexGetTreeParent(dm, point, &parent, NULL));
971     if (parent != point) {
972       PetscInt closureSize, *closure = NULL, i;
973 
974       PetscCall(DMPlexGetTransitiveClosure(dm, parent, PETSC_TRUE, &closureSize, &closure));
975       for (i = 0; i < closureSize; i++) {
976         PetscInt cpoint = closure[2 * i];
977 
978         PetscCall(PetscHSetIAdd(ht, cpoint));
979         PetscCall(DMPlexAddClosure_Tree(dm, ht, cpoint, PETSC_TRUE, PETSC_FALSE));
980       }
981       PetscCall(DMPlexRestoreTransitiveClosure(dm, parent, PETSC_TRUE, &closureSize, &closure));
982     }
983   }
984   if (down) {
985     PetscInt        numChildren;
986     const PetscInt *children;
987 
988     PetscCall(DMPlexGetTreeChildren(dm, point, &numChildren, &children));
989     if (numChildren) {
990       PetscInt i;
991 
992       for (i = 0; i < numChildren; i++) {
993         PetscInt cpoint = children[i];
994 
995         PetscCall(PetscHSetIAdd(ht, cpoint));
996         PetscCall(DMPlexAddClosure_Tree(dm, ht, cpoint, PETSC_FALSE, PETSC_TRUE));
997       }
998     }
999   }
1000   PetscFunctionReturn(PETSC_SUCCESS);
1001 }
1002 
1003 static PetscErrorCode DMPlexAddClosureTree_Up_Private(DM dm, PetscHSetI ht, PetscInt point)
1004 {
1005   PetscInt parent;
1006 
1007   PetscFunctionBeginHot;
1008   PetscCall(DMPlexGetTreeParent(dm, point, &parent, NULL));
1009   if (point != parent) {
1010     const PetscInt *cone;
1011     PetscInt        coneSize, c;
1012 
1013     PetscCall(DMPlexAddClosureTree_Up_Private(dm, ht, parent));
1014     PetscCall(DMPlexAddClosure_Private(dm, ht, parent));
1015     PetscCall(DMPlexGetCone(dm, parent, &cone));
1016     PetscCall(DMPlexGetConeSize(dm, parent, &coneSize));
1017     for (c = 0; c < coneSize; c++) {
1018       const PetscInt cp = cone[c];
1019 
1020       PetscCall(DMPlexAddClosureTree_Up_Private(dm, ht, cp));
1021     }
1022   }
1023   PetscFunctionReturn(PETSC_SUCCESS);
1024 }
1025 
1026 static PetscErrorCode DMPlexAddClosureTree_Down_Private(DM dm, PetscHSetI ht, PetscInt point)
1027 {
1028   PetscInt        i, numChildren;
1029   const PetscInt *children;
1030 
1031   PetscFunctionBeginHot;
1032   PetscCall(DMPlexGetTreeChildren(dm, point, &numChildren, &children));
1033   for (i = 0; i < numChildren; i++) PetscCall(PetscHSetIAdd(ht, children[i]));
1034   PetscFunctionReturn(PETSC_SUCCESS);
1035 }
1036 
1037 static PetscErrorCode DMPlexAddClosureTree_Private(DM dm, PetscHSetI ht, PetscInt point)
1038 {
1039   const PetscInt *cone;
1040   PetscInt        coneSize, c;
1041 
1042   PetscFunctionBeginHot;
1043   PetscCall(PetscHSetIAdd(ht, point));
1044   PetscCall(DMPlexAddClosureTree_Up_Private(dm, ht, point));
1045   PetscCall(DMPlexAddClosureTree_Down_Private(dm, ht, point));
1046   PetscCall(DMPlexGetCone(dm, point, &cone));
1047   PetscCall(DMPlexGetConeSize(dm, point, &coneSize));
1048   for (c = 0; c < coneSize; c++) PetscCall(DMPlexAddClosureTree_Private(dm, ht, cone[c]));
1049   PetscFunctionReturn(PETSC_SUCCESS);
1050 }
1051 
1052 PetscErrorCode DMPlexClosurePoints_Private(DM dm, PetscInt numPoints, const PetscInt points[], IS *closureIS)
1053 {
1054   DM_Plex        *mesh    = (DM_Plex *)dm->data;
1055   const PetscBool hasTree = (mesh->parentSection || mesh->childSection) ? PETSC_TRUE : PETSC_FALSE;
1056   PetscInt        nelems, *elems, off = 0, p;
1057   PetscHSetI      ht = NULL;
1058 
1059   PetscFunctionBegin;
1060   PetscCall(PetscHSetICreate(&ht));
1061   PetscCall(PetscHSetIResize(ht, numPoints * 16));
1062   if (!hasTree) {
1063     for (p = 0; p < numPoints; ++p) PetscCall(DMPlexAddClosure_Private(dm, ht, points[p]));
1064   } else {
1065 #if 1
1066     for (p = 0; p < numPoints; ++p) PetscCall(DMPlexAddClosureTree_Private(dm, ht, points[p]));
1067 #else
1068     PetscInt *closure = NULL, closureSize, c;
1069     for (p = 0; p < numPoints; ++p) {
1070       PetscCall(DMPlexGetTransitiveClosure(dm, points[p], PETSC_TRUE, &closureSize, &closure));
1071       for (c = 0; c < closureSize * 2; c += 2) {
1072         PetscCall(PetscHSetIAdd(ht, closure[c]));
1073         if (hasTree) PetscCall(DMPlexAddClosure_Tree(dm, ht, closure[c], PETSC_TRUE, PETSC_TRUE));
1074       }
1075     }
1076     if (closure) PetscCall(DMPlexRestoreTransitiveClosure(dm, 0, PETSC_TRUE, NULL, &closure));
1077 #endif
1078   }
1079   PetscCall(PetscHSetIGetSize(ht, &nelems));
1080   PetscCall(PetscMalloc1(nelems, &elems));
1081   PetscCall(PetscHSetIGetElems(ht, &off, elems));
1082   PetscCall(PetscHSetIDestroy(&ht));
1083   PetscCall(PetscSortInt(nelems, elems));
1084   PetscCall(ISCreateGeneral(PETSC_COMM_SELF, nelems, elems, PETSC_OWN_POINTER, closureIS));
1085   PetscFunctionReturn(PETSC_SUCCESS);
1086 }
1087 
1088 /*@
1089   DMPlexPartitionLabelClosure - Add the closure of all points to the partition label
1090 
1091   Input Parameters:
1092 + dm    - The `DM`
1093 - label - `DMLabel` assigning ranks to remote roots
1094 
1095   Level: developer
1096 
1097 .seealso: [](ch_unstructured), `DM`, `DMPLEX`, `DMLabel`, `DMPlexPartitionLabelCreateSF()`, `DMPlexDistribute()`, `DMPlexCreateOverlap()`
1098 @*/
1099 PetscErrorCode DMPlexPartitionLabelClosure(DM dm, DMLabel label)
1100 {
1101   IS              rankIS, pointIS, closureIS;
1102   const PetscInt *ranks, *points;
1103   PetscInt        numRanks, numPoints, r;
1104 
1105   PetscFunctionBegin;
1106   PetscCall(DMLabelGetValueIS(label, &rankIS));
1107   PetscCall(ISGetLocalSize(rankIS, &numRanks));
1108   PetscCall(ISGetIndices(rankIS, &ranks));
1109   for (r = 0; r < numRanks; ++r) {
1110     const PetscInt rank = ranks[r];
1111     PetscCall(DMLabelGetStratumIS(label, rank, &pointIS));
1112     PetscCall(ISGetLocalSize(pointIS, &numPoints));
1113     PetscCall(ISGetIndices(pointIS, &points));
1114     PetscCall(DMPlexClosurePoints_Private(dm, numPoints, points, &closureIS));
1115     PetscCall(ISRestoreIndices(pointIS, &points));
1116     PetscCall(ISDestroy(&pointIS));
1117     PetscCall(DMLabelSetStratumIS(label, rank, closureIS));
1118     PetscCall(ISDestroy(&closureIS));
1119   }
1120   PetscCall(ISRestoreIndices(rankIS, &ranks));
1121   PetscCall(ISDestroy(&rankIS));
1122   PetscFunctionReturn(PETSC_SUCCESS);
1123 }
1124 
1125 /*@
1126   DMPlexPartitionLabelAdjacency - Add one level of adjacent points to the partition label
1127 
1128   Input Parameters:
1129 + dm    - The `DM`
1130 - label - `DMLabel` assigning ranks to remote roots
1131 
1132   Level: developer
1133 
1134 .seealso: [](ch_unstructured), `DM`, `DMPLEX`, `DMLabel`, `DMPlexPartitionLabelCreateSF()`, `DMPlexDistribute()`, `DMPlexCreateOverlap()`
1135 @*/
1136 PetscErrorCode DMPlexPartitionLabelAdjacency(DM dm, DMLabel label)
1137 {
1138   IS              rankIS, pointIS;
1139   const PetscInt *ranks, *points;
1140   PetscInt        numRanks, numPoints, r, p, a, adjSize;
1141   PetscInt       *adj = NULL;
1142 
1143   PetscFunctionBegin;
1144   PetscCall(DMLabelGetValueIS(label, &rankIS));
1145   PetscCall(ISGetLocalSize(rankIS, &numRanks));
1146   PetscCall(ISGetIndices(rankIS, &ranks));
1147   for (r = 0; r < numRanks; ++r) {
1148     const PetscInt rank = ranks[r];
1149 
1150     PetscCall(DMLabelGetStratumIS(label, rank, &pointIS));
1151     PetscCall(ISGetLocalSize(pointIS, &numPoints));
1152     PetscCall(ISGetIndices(pointIS, &points));
1153     for (p = 0; p < numPoints; ++p) {
1154       adjSize = PETSC_DETERMINE;
1155       PetscCall(DMPlexGetAdjacency(dm, points[p], &adjSize, &adj));
1156       for (a = 0; a < adjSize; ++a) PetscCall(DMLabelSetValue(label, adj[a], rank));
1157     }
1158     PetscCall(ISRestoreIndices(pointIS, &points));
1159     PetscCall(ISDestroy(&pointIS));
1160   }
1161   PetscCall(ISRestoreIndices(rankIS, &ranks));
1162   PetscCall(ISDestroy(&rankIS));
1163   PetscCall(PetscFree(adj));
1164   PetscFunctionReturn(PETSC_SUCCESS);
1165 }
1166 
1167 /*@
1168   DMPlexPartitionLabelPropagate - Propagate points in a partition label over the point `PetscSF`
1169 
1170   Input Parameters:
1171 + dm    - The `DM`
1172 - label - `DMLabel` assigning ranks to remote roots
1173 
1174   Level: developer
1175 
1176   Note:
1177   This is required when generating multi-level overlaps to capture
1178   overlap points from non-neighbouring partitions.
1179 
1180 .seealso: [](ch_unstructured), `DM`, `DMPLEX`, `DMLabel`, `DMPlexPartitionLabelCreateSF()`, `DMPlexDistribute()`, `DMPlexCreateOverlap()`
1181 @*/
1182 PetscErrorCode DMPlexPartitionLabelPropagate(DM dm, DMLabel label)
1183 {
1184   MPI_Comm        comm;
1185   PetscMPIInt     rank;
1186   PetscSF         sfPoint;
1187   DMLabel         lblRoots, lblLeaves;
1188   IS              rankIS, pointIS;
1189   const PetscInt *ranks;
1190   PetscInt        numRanks, r;
1191 
1192   PetscFunctionBegin;
1193   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
1194   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1195   PetscCall(DMGetPointSF(dm, &sfPoint));
1196   /* Pull point contributions from remote leaves into local roots */
1197   PetscCall(DMLabelGather(label, sfPoint, &lblLeaves));
1198   PetscCall(DMLabelGetValueIS(lblLeaves, &rankIS));
1199   PetscCall(ISGetLocalSize(rankIS, &numRanks));
1200   PetscCall(ISGetIndices(rankIS, &ranks));
1201   for (r = 0; r < numRanks; ++r) {
1202     const PetscInt remoteRank = ranks[r];
1203     if (remoteRank == rank) continue;
1204     PetscCall(DMLabelGetStratumIS(lblLeaves, remoteRank, &pointIS));
1205     PetscCall(DMLabelInsertIS(label, pointIS, remoteRank));
1206     PetscCall(ISDestroy(&pointIS));
1207   }
1208   PetscCall(ISRestoreIndices(rankIS, &ranks));
1209   PetscCall(ISDestroy(&rankIS));
1210   PetscCall(DMLabelDestroy(&lblLeaves));
1211   /* Push point contributions from roots into remote leaves */
1212   PetscCall(DMLabelDistribute(label, sfPoint, &lblRoots));
1213   PetscCall(DMLabelGetValueIS(lblRoots, &rankIS));
1214   PetscCall(ISGetLocalSize(rankIS, &numRanks));
1215   PetscCall(ISGetIndices(rankIS, &ranks));
1216   for (r = 0; r < numRanks; ++r) {
1217     const PetscInt remoteRank = ranks[r];
1218     if (remoteRank == rank) continue;
1219     PetscCall(DMLabelGetStratumIS(lblRoots, remoteRank, &pointIS));
1220     PetscCall(DMLabelInsertIS(label, pointIS, remoteRank));
1221     PetscCall(ISDestroy(&pointIS));
1222   }
1223   PetscCall(ISRestoreIndices(rankIS, &ranks));
1224   PetscCall(ISDestroy(&rankIS));
1225   PetscCall(DMLabelDestroy(&lblRoots));
1226   PetscFunctionReturn(PETSC_SUCCESS);
1227 }
1228 
1229 /*@
1230   DMPlexPartitionLabelInvert - Create a partition label of remote roots from a local root label
1231 
1232   Input Parameters:
1233 + dm        - The `DM`
1234 . rootLabel - `DMLabel` assigning ranks to local roots
1235 - processSF - A star forest mapping into the local index on each remote rank
1236 
1237   Output Parameter:
1238 . leafLabel - `DMLabel` assigning ranks to remote roots
1239 
1240   Level: developer
1241 
1242   Note:
1243   The rootLabel defines a send pattern by mapping local points to remote target ranks. The
1244   resulting leafLabel is a receiver mapping of remote roots to their parent rank.
1245 
1246 .seealso: [](ch_unstructured), `DM`, `DMPLEX`, `DMPlexPartitionLabelCreateSF()`, `DMPlexDistribute()`, `DMPlexCreateOverlap()`
1247 @*/
1248 PetscErrorCode DMPlexPartitionLabelInvert(DM dm, DMLabel rootLabel, PetscSF processSF, DMLabel leafLabel)
1249 {
1250   MPI_Comm           comm;
1251   PetscMPIInt        rank, size, r;
1252   PetscInt           p, n, numNeighbors, numPoints, dof, off, rootSize, l, nleaves, leafSize;
1253   PetscSF            sfPoint;
1254   PetscSection       rootSection;
1255   PetscSFNode       *rootPoints, *leafPoints;
1256   const PetscSFNode *remote;
1257   const PetscInt    *local, *neighbors;
1258   IS                 valueIS;
1259   PetscBool          mpiOverflow = PETSC_FALSE;
1260 
1261   PetscFunctionBegin;
1262   PetscCall(PetscLogEventBegin(DMPLEX_PartLabelInvert, dm, 0, 0, 0));
1263   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
1264   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1265   PetscCallMPI(MPI_Comm_size(comm, &size));
1266   PetscCall(DMGetPointSF(dm, &sfPoint));
1267 
1268   /* Convert to (point, rank) and use actual owners */
1269   PetscCall(PetscSectionCreate(comm, &rootSection));
1270   PetscCall(PetscSectionSetChart(rootSection, 0, size));
1271   PetscCall(DMLabelGetValueIS(rootLabel, &valueIS));
1272   PetscCall(ISGetLocalSize(valueIS, &numNeighbors));
1273   PetscCall(ISGetIndices(valueIS, &neighbors));
1274   for (n = 0; n < numNeighbors; ++n) {
1275     PetscCall(DMLabelGetStratumSize(rootLabel, neighbors[n], &numPoints));
1276     PetscCall(PetscSectionAddDof(rootSection, neighbors[n], numPoints));
1277   }
1278   PetscCall(PetscSectionSetUp(rootSection));
1279   PetscCall(PetscSectionGetStorageSize(rootSection, &rootSize));
1280   PetscCall(PetscMalloc1(rootSize, &rootPoints));
1281   PetscCall(PetscSFGetGraph(sfPoint, NULL, &nleaves, &local, &remote));
1282   for (n = 0; n < numNeighbors; ++n) {
1283     IS              pointIS;
1284     const PetscInt *points;
1285 
1286     PetscCall(PetscSectionGetOffset(rootSection, neighbors[n], &off));
1287     PetscCall(DMLabelGetStratumIS(rootLabel, neighbors[n], &pointIS));
1288     PetscCall(ISGetLocalSize(pointIS, &numPoints));
1289     PetscCall(ISGetIndices(pointIS, &points));
1290     for (p = 0; p < numPoints; ++p) {
1291       if (local) PetscCall(PetscFindInt(points[p], nleaves, local, &l));
1292       else l = -1;
1293       if (l >= 0) {
1294         rootPoints[off + p] = remote[l];
1295       } else {
1296         rootPoints[off + p].index = points[p];
1297         rootPoints[off + p].rank  = rank;
1298       }
1299     }
1300     PetscCall(ISRestoreIndices(pointIS, &points));
1301     PetscCall(ISDestroy(&pointIS));
1302   }
1303 
1304   /* Try to communicate overlap using All-to-All */
1305   if (!processSF) {
1306     PetscInt64   counter     = 0;
1307     PetscBool    locOverflow = PETSC_FALSE;
1308     PetscMPIInt *scounts, *sdispls, *rcounts, *rdispls;
1309 
1310     PetscCall(PetscCalloc4(size, &scounts, size, &sdispls, size, &rcounts, size, &rdispls));
1311     for (n = 0; n < numNeighbors; ++n) {
1312       PetscCall(PetscSectionGetDof(rootSection, neighbors[n], &dof));
1313       PetscCall(PetscSectionGetOffset(rootSection, neighbors[n], &off));
1314 #if defined(PETSC_USE_64BIT_INDICES)
1315       if (dof > PETSC_MPI_INT_MAX) {
1316         locOverflow = PETSC_TRUE;
1317         break;
1318       }
1319       if (off > PETSC_MPI_INT_MAX) {
1320         locOverflow = PETSC_TRUE;
1321         break;
1322       }
1323 #endif
1324       scounts[neighbors[n]] = (PetscMPIInt)dof;
1325       sdispls[neighbors[n]] = (PetscMPIInt)off;
1326     }
1327     PetscCallMPI(MPI_Alltoall(scounts, 1, MPI_INT, rcounts, 1, MPI_INT, comm));
1328     for (r = 0; r < size; ++r) {
1329       rdispls[r] = (int)counter;
1330       counter += rcounts[r];
1331     }
1332     if (counter > PETSC_MPI_INT_MAX) locOverflow = PETSC_TRUE;
1333     PetscCall(MPIU_Allreduce(&locOverflow, &mpiOverflow, 1, MPIU_BOOL, MPI_LOR, comm));
1334     if (!mpiOverflow) {
1335       PetscCall(PetscInfo(dm, "Using Alltoallv for mesh distribution\n"));
1336       leafSize = (PetscInt)counter;
1337       PetscCall(PetscMalloc1(leafSize, &leafPoints));
1338       PetscCallMPI(MPI_Alltoallv(rootPoints, scounts, sdispls, MPIU_2INT, leafPoints, rcounts, rdispls, MPIU_2INT, comm));
1339     }
1340     PetscCall(PetscFree4(scounts, sdispls, rcounts, rdispls));
1341   }
1342 
1343   /* Communicate overlap using process star forest */
1344   if (processSF || mpiOverflow) {
1345     PetscSF      procSF;
1346     PetscSection leafSection;
1347 
1348     if (processSF) {
1349       PetscCall(PetscInfo(dm, "Using processSF for mesh distribution\n"));
1350       PetscCall(PetscObjectReference((PetscObject)processSF));
1351       procSF = processSF;
1352     } else {
1353       PetscCall(PetscInfo(dm, "Using processSF for mesh distribution (MPI overflow)\n"));
1354       PetscCall(PetscSFCreate(comm, &procSF));
1355       PetscCall(PetscSFSetGraphWithPattern(procSF, NULL, PETSCSF_PATTERN_ALLTOALL));
1356     }
1357 
1358     PetscCall(PetscSectionCreate(PetscObjectComm((PetscObject)dm), &leafSection));
1359     PetscCall(DMPlexDistributeData(dm, procSF, rootSection, MPIU_2INT, rootPoints, leafSection, (void **)&leafPoints));
1360     PetscCall(PetscSectionGetStorageSize(leafSection, &leafSize));
1361     PetscCall(PetscSectionDestroy(&leafSection));
1362     PetscCall(PetscSFDestroy(&procSF));
1363   }
1364 
1365   for (p = 0; p < leafSize; p++) PetscCall(DMLabelSetValue(leafLabel, leafPoints[p].index, leafPoints[p].rank));
1366 
1367   PetscCall(ISRestoreIndices(valueIS, &neighbors));
1368   PetscCall(ISDestroy(&valueIS));
1369   PetscCall(PetscSectionDestroy(&rootSection));
1370   PetscCall(PetscFree(rootPoints));
1371   PetscCall(PetscFree(leafPoints));
1372   PetscCall(PetscLogEventEnd(DMPLEX_PartLabelInvert, dm, 0, 0, 0));
1373   PetscFunctionReturn(PETSC_SUCCESS);
1374 }
1375 
1376 /*@
1377   DMPlexPartitionLabelCreateSF - Create a star forest from a label that assigns ranks to points
1378 
1379   Input Parameters:
1380 + dm        - The `DM`
1381 . label     - `DMLabel` assigning ranks to remote roots
1382 - sortRanks - Whether or not to sort the SF leaves by rank
1383 
1384   Output Parameter:
1385 . sf - The star forest communication context encapsulating the defined mapping
1386 
1387   Level: developer
1388 
1389   Note:
1390   The incoming label is a receiver mapping of remote points to their parent rank.
1391 
1392 .seealso: [](ch_unstructured), `DM`, `DMPLEX`, `DMLabel`, `PetscSF`, `DMPlexDistribute()`, `DMPlexCreateOverlap()`
1393 @*/
1394 PetscErrorCode DMPlexPartitionLabelCreateSF(DM dm, DMLabel label, PetscBool sortRanks, PetscSF *sf)
1395 {
1396   PetscMPIInt     rank;
1397   PetscInt        n, numRemote, p, numPoints, pStart, pEnd, idx = 0, nNeighbors;
1398   PetscSFNode    *remotePoints;
1399   IS              remoteRootIS, neighborsIS;
1400   const PetscInt *remoteRoots, *neighbors;
1401   PetscMPIInt     myRank;
1402 
1403   PetscFunctionBegin;
1404   PetscCall(PetscLogEventBegin(DMPLEX_PartLabelCreateSF, dm, 0, 0, 0));
1405   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank));
1406   PetscCall(DMLabelGetValueIS(label, &neighborsIS));
1407 
1408   if (sortRanks) {
1409     IS is;
1410 
1411     PetscCall(ISDuplicate(neighborsIS, &is));
1412     PetscCall(ISSort(is));
1413     PetscCall(ISDestroy(&neighborsIS));
1414     neighborsIS = is;
1415   }
1416   myRank = sortRanks ? -1 : rank;
1417 
1418   PetscCall(ISGetLocalSize(neighborsIS, &nNeighbors));
1419   PetscCall(ISGetIndices(neighborsIS, &neighbors));
1420   for (numRemote = 0, n = 0; n < nNeighbors; ++n) {
1421     PetscCall(DMLabelGetStratumSize(label, neighbors[n], &numPoints));
1422     numRemote += numPoints;
1423   }
1424   PetscCall(PetscMalloc1(numRemote, &remotePoints));
1425 
1426   /* Put owned points first if not sorting the ranks */
1427   if (!sortRanks) {
1428     PetscCall(DMLabelGetStratumSize(label, rank, &numPoints));
1429     if (numPoints > 0) {
1430       PetscCall(DMLabelGetStratumIS(label, rank, &remoteRootIS));
1431       PetscCall(ISGetIndices(remoteRootIS, &remoteRoots));
1432       for (p = 0; p < numPoints; p++) {
1433         remotePoints[idx].index = remoteRoots[p];
1434         remotePoints[idx].rank  = rank;
1435         idx++;
1436       }
1437       PetscCall(ISRestoreIndices(remoteRootIS, &remoteRoots));
1438       PetscCall(ISDestroy(&remoteRootIS));
1439     }
1440   }
1441 
1442   /* Now add remote points */
1443   for (n = 0; n < nNeighbors; ++n) {
1444     const PetscInt nn = neighbors[n];
1445 
1446     PetscCall(DMLabelGetStratumSize(label, nn, &numPoints));
1447     if (nn == myRank || numPoints <= 0) continue;
1448     PetscCall(DMLabelGetStratumIS(label, nn, &remoteRootIS));
1449     PetscCall(ISGetIndices(remoteRootIS, &remoteRoots));
1450     for (p = 0; p < numPoints; p++) {
1451       remotePoints[idx].index = remoteRoots[p];
1452       remotePoints[idx].rank  = nn;
1453       idx++;
1454     }
1455     PetscCall(ISRestoreIndices(remoteRootIS, &remoteRoots));
1456     PetscCall(ISDestroy(&remoteRootIS));
1457   }
1458 
1459   PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)dm), sf));
1460   PetscCall(DMPlexGetChart(dm, &pStart, &pEnd));
1461   PetscCall(PetscSFSetGraph(*sf, pEnd - pStart, numRemote, NULL, PETSC_OWN_POINTER, remotePoints, PETSC_OWN_POINTER));
1462   PetscCall(PetscSFSetUp(*sf));
1463   PetscCall(ISDestroy(&neighborsIS));
1464   PetscCall(PetscLogEventEnd(DMPLEX_PartLabelCreateSF, dm, 0, 0, 0));
1465   PetscFunctionReturn(PETSC_SUCCESS);
1466 }
1467 
1468 #if defined(PETSC_HAVE_PARMETIS)
1469   #include <parmetis.h>
1470 #endif
1471 
1472 /* The two functions below are used by DMPlexRebalanceSharedPoints which errors
1473  * when PETSc is built without ParMETIS. To avoid -Wunused-function, we take
1474  * them out in that case. */
1475 #if defined(PETSC_HAVE_PARMETIS)
1476 /*
1477   DMPlexRewriteSF - Rewrites the ownership of the `PetsSF` of a `DM` (in place).
1478 
1479   Input parameters:
1480 + dm                - The `DMPLEX` object.
1481 . n                 - The number of points.
1482 . pointsToRewrite   - The points in the `PetscSF` whose ownership will change.
1483 . targetOwners      - New owner for each element in pointsToRewrite.
1484 - degrees           - Degrees of the points in the `PetscSF` as obtained by `PetscSFComputeDegreeBegin()`/`PetscSFComputeDegreeEnd()`.
1485 
1486   Level: developer
1487 
1488 .seealso: [](ch_unstructured), `DM`, `DMPLEX`, `DMLabel`, `PetscSF`, `DMPlexDistribute()`, `DMPlexCreateOverlap()`
1489 */
1490 static PetscErrorCode DMPlexRewriteSF(DM dm, PetscInt n, PetscInt *pointsToRewrite, PetscInt *targetOwners, const PetscInt *degrees)
1491 {
1492   PetscInt           pStart, pEnd, i, j, counter, leafCounter, sumDegrees, nroots, nleafs;
1493   PetscInt          *cumSumDegrees, *newOwners, *newNumbers, *rankOnLeafs, *locationsOfLeafs, *remoteLocalPointOfLeafs, *points, *leafsNew;
1494   PetscSFNode       *leafLocationsNew;
1495   const PetscSFNode *iremote;
1496   const PetscInt    *ilocal;
1497   PetscBool         *isLeaf;
1498   PetscSF            sf;
1499   MPI_Comm           comm;
1500   PetscMPIInt        rank, size;
1501 
1502   PetscFunctionBegin;
1503   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
1504   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1505   PetscCallMPI(MPI_Comm_size(comm, &size));
1506   PetscCall(DMPlexGetChart(dm, &pStart, &pEnd));
1507 
1508   PetscCall(DMGetPointSF(dm, &sf));
1509   PetscCall(PetscSFGetGraph(sf, &nroots, &nleafs, &ilocal, &iremote));
1510   PetscCall(PetscMalloc1(pEnd - pStart, &isLeaf));
1511   for (i = 0; i < pEnd - pStart; i++) isLeaf[i] = PETSC_FALSE;
1512   for (i = 0; i < nleafs; i++) isLeaf[ilocal[i] - pStart] = PETSC_TRUE;
1513 
1514   PetscCall(PetscMalloc1(pEnd - pStart + 1, &cumSumDegrees));
1515   cumSumDegrees[0] = 0;
1516   for (i = 1; i <= pEnd - pStart; i++) cumSumDegrees[i] = cumSumDegrees[i - 1] + degrees[i - 1];
1517   sumDegrees = cumSumDegrees[pEnd - pStart];
1518   /* get the location of my leafs (we have sumDegrees many leafs pointing at our roots) */
1519 
1520   PetscCall(PetscMalloc1(sumDegrees, &locationsOfLeafs));
1521   PetscCall(PetscMalloc1(pEnd - pStart, &rankOnLeafs));
1522   for (i = 0; i < pEnd - pStart; i++) rankOnLeafs[i] = rank;
1523   PetscCall(PetscSFGatherBegin(sf, MPIU_INT, rankOnLeafs, locationsOfLeafs));
1524   PetscCall(PetscSFGatherEnd(sf, MPIU_INT, rankOnLeafs, locationsOfLeafs));
1525   PetscCall(PetscFree(rankOnLeafs));
1526 
1527   /* get the remote local points of my leaves */
1528   PetscCall(PetscMalloc1(sumDegrees, &remoteLocalPointOfLeafs));
1529   PetscCall(PetscMalloc1(pEnd - pStart, &points));
1530   for (i = 0; i < pEnd - pStart; i++) points[i] = pStart + i;
1531   PetscCall(PetscSFGatherBegin(sf, MPIU_INT, points, remoteLocalPointOfLeafs));
1532   PetscCall(PetscSFGatherEnd(sf, MPIU_INT, points, remoteLocalPointOfLeafs));
1533   PetscCall(PetscFree(points));
1534   /* Figure out the new owners of the vertices that are up for grabs and their numbers on the new owners */
1535   PetscCall(PetscMalloc1(pEnd - pStart, &newOwners));
1536   PetscCall(PetscMalloc1(pEnd - pStart, &newNumbers));
1537   for (i = 0; i < pEnd - pStart; i++) {
1538     newOwners[i]  = -1;
1539     newNumbers[i] = -1;
1540   }
1541   {
1542     PetscInt oldNumber, newNumber, oldOwner, newOwner;
1543     for (i = 0; i < n; i++) {
1544       oldNumber = pointsToRewrite[i];
1545       newNumber = -1;
1546       oldOwner  = rank;
1547       newOwner  = targetOwners[i];
1548       if (oldOwner == newOwner) {
1549         newNumber = oldNumber;
1550       } else {
1551         for (j = 0; j < degrees[oldNumber]; j++) {
1552           if (locationsOfLeafs[cumSumDegrees[oldNumber] + j] == newOwner) {
1553             newNumber = remoteLocalPointOfLeafs[cumSumDegrees[oldNumber] + j];
1554             break;
1555           }
1556         }
1557       }
1558       PetscCheck(newNumber != -1, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Couldn't find the new owner of vertex.");
1559 
1560       newOwners[oldNumber]  = newOwner;
1561       newNumbers[oldNumber] = newNumber;
1562     }
1563   }
1564   PetscCall(PetscFree(cumSumDegrees));
1565   PetscCall(PetscFree(locationsOfLeafs));
1566   PetscCall(PetscFree(remoteLocalPointOfLeafs));
1567 
1568   PetscCall(PetscSFBcastBegin(sf, MPIU_INT, newOwners, newOwners, MPI_REPLACE));
1569   PetscCall(PetscSFBcastEnd(sf, MPIU_INT, newOwners, newOwners, MPI_REPLACE));
1570   PetscCall(PetscSFBcastBegin(sf, MPIU_INT, newNumbers, newNumbers, MPI_REPLACE));
1571   PetscCall(PetscSFBcastEnd(sf, MPIU_INT, newNumbers, newNumbers, MPI_REPLACE));
1572 
1573   /* Now count how many leafs we have on each processor. */
1574   leafCounter = 0;
1575   for (i = 0; i < pEnd - pStart; i++) {
1576     if (newOwners[i] >= 0) {
1577       if (newOwners[i] != rank) leafCounter++;
1578     } else {
1579       if (isLeaf[i]) leafCounter++;
1580     }
1581   }
1582 
1583   /* Now set up the new sf by creating the leaf arrays */
1584   PetscCall(PetscMalloc1(leafCounter, &leafsNew));
1585   PetscCall(PetscMalloc1(leafCounter, &leafLocationsNew));
1586 
1587   leafCounter = 0;
1588   counter     = 0;
1589   for (i = 0; i < pEnd - pStart; i++) {
1590     if (newOwners[i] >= 0) {
1591       if (newOwners[i] != rank) {
1592         leafsNew[leafCounter]               = i;
1593         leafLocationsNew[leafCounter].rank  = newOwners[i];
1594         leafLocationsNew[leafCounter].index = newNumbers[i];
1595         leafCounter++;
1596       }
1597     } else {
1598       if (isLeaf[i]) {
1599         leafsNew[leafCounter]               = i;
1600         leafLocationsNew[leafCounter].rank  = iremote[counter].rank;
1601         leafLocationsNew[leafCounter].index = iremote[counter].index;
1602         leafCounter++;
1603       }
1604     }
1605     if (isLeaf[i]) counter++;
1606   }
1607 
1608   PetscCall(PetscSFSetGraph(sf, nroots, leafCounter, leafsNew, PETSC_OWN_POINTER, leafLocationsNew, PETSC_OWN_POINTER));
1609   PetscCall(PetscFree(newOwners));
1610   PetscCall(PetscFree(newNumbers));
1611   PetscCall(PetscFree(isLeaf));
1612   PetscFunctionReturn(PETSC_SUCCESS);
1613 }
1614 
1615 static PetscErrorCode DMPlexViewDistribution(MPI_Comm comm, PetscInt n, PetscInt skip, PetscInt *vtxwgt, PetscInt *part, PetscViewer viewer)
1616 {
1617   PetscInt   *distribution, min, max, sum;
1618   PetscMPIInt rank, size;
1619 
1620   PetscFunctionBegin;
1621   PetscCallMPI(MPI_Comm_size(comm, &size));
1622   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1623   PetscCall(PetscCalloc1(size, &distribution));
1624   for (PetscInt i = 0; i < n; i++) {
1625     if (part) distribution[part[i]] += vtxwgt[skip * i];
1626     else distribution[rank] += vtxwgt[skip * i];
1627   }
1628   PetscCall(MPIU_Allreduce(MPI_IN_PLACE, distribution, size, MPIU_INT, MPI_SUM, comm));
1629   min = distribution[0];
1630   max = distribution[0];
1631   sum = distribution[0];
1632   for (PetscInt i = 1; i < size; i++) {
1633     if (distribution[i] < min) min = distribution[i];
1634     if (distribution[i] > max) max = distribution[i];
1635     sum += distribution[i];
1636   }
1637   PetscCall(PetscViewerASCIIPrintf(viewer, "Min: %" PetscInt_FMT ", Avg: %" PetscInt_FMT ", Max: %" PetscInt_FMT ", Balance: %f\n", min, sum / size, max, (max * 1. * size) / sum));
1638   PetscCall(PetscFree(distribution));
1639   PetscFunctionReturn(PETSC_SUCCESS);
1640 }
1641 
1642 #endif
1643 
1644 /*@
1645   DMPlexRebalanceSharedPoints - Redistribute points in the plex that are shared in order to achieve better balancing. This routine updates the `PointSF` of the `DM` inplace.
1646 
1647   Input Parameters:
1648 + dm              - The `DMPLEX` object.
1649 . entityDepth     - depth of the entity to balance (0 -> balance vertices).
1650 . useInitialGuess - whether to use the current distribution as initial guess (only used by ParMETIS).
1651 - parallel        - whether to use ParMETIS and do the partition in parallel or whether to gather the graph onto a single process and use METIS.
1652 
1653   Output Parameter:
1654 . success - whether the graph partitioning was successful or not, optional. Unsuccessful simply means no change to the partitioning
1655 
1656   Options Database Keys:
1657 + -dm_plex_rebalance_shared_points_parmetis             - Use ParMetis instead of Metis for the partitioner
1658 . -dm_plex_rebalance_shared_points_use_initial_guess    - Use current partition to bootstrap ParMetis partition
1659 . -dm_plex_rebalance_shared_points_use_mat_partitioning - Use the MatPartitioning object to perform the partition, the prefix for those operations is -dm_plex_rebalance_shared_points_
1660 - -dm_plex_rebalance_shared_points_monitor              - Monitor the shared points rebalance process
1661 
1662   Level: intermediate
1663 
1664 .seealso: [](ch_unstructured), `DM`, `DMPLEX`, `DMPlexDistribute()`, `DMPlexCreateOverlap()`
1665 @*/
1666 PetscErrorCode DMPlexRebalanceSharedPoints(DM dm, PetscInt entityDepth, PetscBool useInitialGuess, PetscBool parallel, PetscBool *success)
1667 {
1668 #if defined(PETSC_HAVE_PARMETIS)
1669   PetscSF            sf;
1670   PetscInt           ierr, i, j, idx, jdx;
1671   PetscInt           eBegin, eEnd, nroots, nleafs, pStart, pEnd;
1672   const PetscInt    *degrees, *ilocal;
1673   const PetscSFNode *iremote;
1674   PetscBool         *toBalance, *isLeaf, *isExclusivelyOwned, *isNonExclusivelyOwned;
1675   PetscInt           numExclusivelyOwned, numNonExclusivelyOwned;
1676   PetscMPIInt        rank, size;
1677   PetscInt          *globalNumbersOfLocalOwnedVertices, *leafGlobalNumbers;
1678   const PetscInt    *cumSumVertices;
1679   PetscInt           offset, counter;
1680   PetscInt          *vtxwgt;
1681   const PetscInt    *xadj, *adjncy;
1682   PetscInt          *part, *options;
1683   PetscInt           nparts, wgtflag, numflag, ncon, edgecut;
1684   real_t            *ubvec;
1685   PetscInt          *firstVertices, *renumbering;
1686   PetscInt           failed, failedGlobal;
1687   MPI_Comm           comm;
1688   Mat                A;
1689   PetscViewer        viewer;
1690   PetscViewerFormat  format;
1691   PetscLayout        layout;
1692   real_t            *tpwgts;
1693   PetscMPIInt       *counts, *mpiCumSumVertices;
1694   PetscInt          *pointsToRewrite;
1695   PetscInt           numRows;
1696   PetscBool          done, usematpartitioning = PETSC_FALSE;
1697   IS                 ispart = NULL;
1698   MatPartitioning    mp;
1699   const char        *prefix;
1700 
1701   PetscFunctionBegin;
1702   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
1703   PetscCallMPI(MPI_Comm_size(comm, &size));
1704   if (size == 1) {
1705     if (success) *success = PETSC_TRUE;
1706     PetscFunctionReturn(PETSC_SUCCESS);
1707   }
1708   if (success) *success = PETSC_FALSE;
1709   PetscCallMPI(MPI_Comm_rank(comm, &rank));
1710 
1711   parallel        = PETSC_FALSE;
1712   useInitialGuess = PETSC_FALSE;
1713   PetscObjectOptionsBegin((PetscObject)dm);
1714   PetscCall(PetscOptionsName("-dm_plex_rebalance_shared_points_parmetis", "Use ParMetis instead of Metis for the partitioner", "DMPlexRebalanceSharedPoints", &parallel));
1715   PetscCall(PetscOptionsBool("-dm_plex_rebalance_shared_points_use_initial_guess", "Use current partition to bootstrap ParMetis partition", "DMPlexRebalanceSharedPoints", useInitialGuess, &useInitialGuess, NULL));
1716   PetscCall(PetscOptionsBool("-dm_plex_rebalance_shared_points_use_mat_partitioning", "Use the MatPartitioning object to partition", "DMPlexRebalanceSharedPoints", usematpartitioning, &usematpartitioning, NULL));
1717   PetscCall(PetscOptionsViewer("-dm_plex_rebalance_shared_points_monitor", "Monitor the shared points rebalance process", "DMPlexRebalanceSharedPoints", &viewer, &format, NULL));
1718   PetscOptionsEnd();
1719   if (viewer) PetscCall(PetscViewerPushFormat(viewer, format));
1720 
1721   PetscCall(PetscLogEventBegin(DMPLEX_RebalanceSharedPoints, dm, 0, 0, 0));
1722 
1723   PetscCall(DMGetOptionsPrefix(dm, &prefix));
1724   PetscCall(PetscOptionsGetViewer(comm, ((PetscObject)dm)->options, prefix, "-dm_rebalance_partition_view", &viewer, &format, NULL));
1725   if (viewer) PetscCall(PetscViewerPushFormat(viewer, format));
1726 
1727   /* Figure out all points in the plex that we are interested in balancing. */
1728   PetscCall(DMPlexGetDepthStratum(dm, entityDepth, &eBegin, &eEnd));
1729   PetscCall(DMPlexGetChart(dm, &pStart, &pEnd));
1730   PetscCall(PetscMalloc1(pEnd - pStart, &toBalance));
1731   for (i = 0; i < pEnd - pStart; i++) toBalance[i] = (PetscBool)(i >= eBegin && i < eEnd);
1732 
1733   /* There are three types of points:
1734    * exclusivelyOwned: points that are owned by this process and only seen by this process
1735    * nonExclusivelyOwned: points that are owned by this process but seen by at least another process
1736    * leaf: a point that is seen by this process but owned by a different process
1737    */
1738   PetscCall(DMGetPointSF(dm, &sf));
1739   PetscCall(PetscSFGetGraph(sf, &nroots, &nleafs, &ilocal, &iremote));
1740   PetscCall(PetscMalloc1(pEnd - pStart, &isLeaf));
1741   PetscCall(PetscMalloc1(pEnd - pStart, &isNonExclusivelyOwned));
1742   PetscCall(PetscMalloc1(pEnd - pStart, &isExclusivelyOwned));
1743   for (i = 0; i < pEnd - pStart; i++) {
1744     isNonExclusivelyOwned[i] = PETSC_FALSE;
1745     isExclusivelyOwned[i]    = PETSC_FALSE;
1746     isLeaf[i]                = PETSC_FALSE;
1747   }
1748 
1749   /* mark all the leafs */
1750   for (i = 0; i < nleafs; i++) isLeaf[ilocal[i] - pStart] = PETSC_TRUE;
1751 
1752   /* for an owned point, we can figure out whether another processor sees it or
1753    * not by calculating its degree */
1754   PetscCall(PetscSFComputeDegreeBegin(sf, &degrees));
1755   PetscCall(PetscSFComputeDegreeEnd(sf, &degrees));
1756   numExclusivelyOwned    = 0;
1757   numNonExclusivelyOwned = 0;
1758   for (i = 0; i < pEnd - pStart; i++) {
1759     if (toBalance[i]) {
1760       if (degrees[i] > 0) {
1761         isNonExclusivelyOwned[i] = PETSC_TRUE;
1762         numNonExclusivelyOwned += 1;
1763       } else {
1764         if (!isLeaf[i]) {
1765           isExclusivelyOwned[i] = PETSC_TRUE;
1766           numExclusivelyOwned += 1;
1767         }
1768       }
1769     }
1770   }
1771 
1772   /* Build a graph with one vertex per core representing the
1773    * exclusively owned points and then one vertex per nonExclusively owned
1774    * point. */
1775   PetscCall(PetscLayoutCreate(comm, &layout));
1776   PetscCall(PetscLayoutSetLocalSize(layout, 1 + numNonExclusivelyOwned));
1777   PetscCall(PetscLayoutSetUp(layout));
1778   PetscCall(PetscLayoutGetRanges(layout, &cumSumVertices));
1779   PetscCall(PetscMalloc1(pEnd - pStart, &globalNumbersOfLocalOwnedVertices));
1780   for (i = 0; i < pEnd - pStart; i++) globalNumbersOfLocalOwnedVertices[i] = pStart - 1;
1781   offset  = cumSumVertices[rank];
1782   counter = 0;
1783   for (i = 0; i < pEnd - pStart; i++) {
1784     if (toBalance[i]) {
1785       if (degrees[i] > 0) {
1786         globalNumbersOfLocalOwnedVertices[i] = counter + 1 + offset;
1787         counter++;
1788       }
1789     }
1790   }
1791 
1792   /* send the global numbers of vertices I own to the leafs so that they know to connect to it */
1793   PetscCall(PetscMalloc1(pEnd - pStart, &leafGlobalNumbers));
1794   PetscCall(PetscSFBcastBegin(sf, MPIU_INT, globalNumbersOfLocalOwnedVertices, leafGlobalNumbers, MPI_REPLACE));
1795   PetscCall(PetscSFBcastEnd(sf, MPIU_INT, globalNumbersOfLocalOwnedVertices, leafGlobalNumbers, MPI_REPLACE));
1796 
1797   /* Build the graph for partitioning */
1798   numRows = 1 + numNonExclusivelyOwned;
1799   PetscCall(PetscLogEventBegin(DMPLEX_RebalBuildGraph, dm, 0, 0, 0));
1800   PetscCall(MatCreate(comm, &A));
1801   PetscCall(MatSetType(A, MATMPIADJ));
1802   PetscCall(MatSetSizes(A, numRows, numRows, cumSumVertices[size], cumSumVertices[size]));
1803   idx = cumSumVertices[rank];
1804   for (i = 0; i < pEnd - pStart; i++) {
1805     if (toBalance[i]) {
1806       if (isNonExclusivelyOwned[i]) jdx = globalNumbersOfLocalOwnedVertices[i];
1807       else if (isLeaf[i]) jdx = leafGlobalNumbers[i];
1808       else continue;
1809       PetscCall(MatSetValue(A, idx, jdx, 1, INSERT_VALUES));
1810       PetscCall(MatSetValue(A, jdx, idx, 1, INSERT_VALUES));
1811     }
1812   }
1813   PetscCall(PetscFree(globalNumbersOfLocalOwnedVertices));
1814   PetscCall(PetscFree(leafGlobalNumbers));
1815   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
1816   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
1817   PetscCall(PetscLogEventEnd(DMPLEX_RebalBuildGraph, dm, 0, 0, 0));
1818 
1819   nparts = size;
1820   ncon   = 1;
1821   PetscCall(PetscMalloc1(ncon * nparts, &tpwgts));
1822   for (i = 0; i < ncon * nparts; i++) tpwgts[i] = 1. / (nparts);
1823   PetscCall(PetscMalloc1(ncon, &ubvec));
1824   for (i = 0; i < ncon; i++) ubvec[i] = 1.05;
1825 
1826   PetscCall(PetscMalloc1(ncon * (1 + numNonExclusivelyOwned), &vtxwgt));
1827   if (ncon == 2) {
1828     vtxwgt[0] = numExclusivelyOwned;
1829     vtxwgt[1] = 1;
1830     for (i = 0; i < numNonExclusivelyOwned; i++) {
1831       vtxwgt[ncon * (i + 1)]     = 1;
1832       vtxwgt[ncon * (i + 1) + 1] = 0;
1833     }
1834   } else {
1835     PetscInt base, ms;
1836     PetscCall(MPIU_Allreduce(&numExclusivelyOwned, &base, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)dm)));
1837     PetscCall(MatGetSize(A, &ms, NULL));
1838     ms -= size;
1839     base      = PetscMax(base, ms);
1840     vtxwgt[0] = base + numExclusivelyOwned;
1841     for (i = 0; i < numNonExclusivelyOwned; i++) vtxwgt[i + 1] = 1;
1842   }
1843 
1844   if (viewer) {
1845     PetscCall(PetscViewerASCIIPrintf(viewer, "Attempt rebalancing of shared points of depth %" PetscInt_FMT " on interface of mesh distribution.\n", entityDepth));
1846     PetscCall(PetscViewerASCIIPrintf(viewer, "Size of generated auxiliary graph: %" PetscInt_FMT "\n", cumSumVertices[size]));
1847   }
1848   /* TODO: Drop the parallel/sequential choice here and just use MatPartioner for much more flexibility */
1849   if (usematpartitioning) {
1850     const char *prefix;
1851 
1852     PetscCall(MatPartitioningCreate(PetscObjectComm((PetscObject)dm), &mp));
1853     PetscCall(PetscObjectSetOptionsPrefix((PetscObject)mp, "dm_plex_rebalance_shared_points_"));
1854     PetscCall(PetscObjectGetOptionsPrefix((PetscObject)dm, &prefix));
1855     PetscCall(PetscObjectPrependOptionsPrefix((PetscObject)mp, prefix));
1856     PetscCall(MatPartitioningSetAdjacency(mp, A));
1857     PetscCall(MatPartitioningSetNumberVertexWeights(mp, ncon));
1858     PetscCall(MatPartitioningSetVertexWeights(mp, vtxwgt));
1859     PetscCall(MatPartitioningSetFromOptions(mp));
1860     PetscCall(MatPartitioningApply(mp, &ispart));
1861     PetscCall(ISGetIndices(ispart, (const PetscInt **)&part));
1862   } else if (parallel) {
1863     if (viewer) PetscCall(PetscViewerASCIIPrintf(viewer, "Using ParMETIS to partition graph.\n"));
1864     PetscCall(PetscMalloc1(cumSumVertices[rank + 1] - cumSumVertices[rank], &part));
1865     PetscCall(MatGetRowIJ(A, PETSC_FALSE, PETSC_FALSE, PETSC_FALSE, &numRows, &xadj, &adjncy, &done));
1866     PetscCheck(done, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Could not get adjacency information");
1867     PetscCall(PetscMalloc1(4, &options));
1868     options[0] = 1;
1869     options[1] = 0; /* Verbosity */
1870     if (viewer) options[1] = 1;
1871     options[2] = 0;                    /* Seed */
1872     options[3] = PARMETIS_PSR_COUPLED; /* Seed */
1873     wgtflag    = 2;
1874     numflag    = 0;
1875     if (useInitialGuess) {
1876       PetscCall(PetscViewerASCIIPrintf(viewer, "THIS DOES NOT WORK! I don't know why. Using current distribution of points as initial guess.\n"));
1877       for (i = 0; i < numRows; i++) part[i] = rank;
1878       if (viewer) PetscCall(PetscViewerASCIIPrintf(viewer, "Using current distribution of points as initial guess.\n"));
1879       PetscStackPushExternal("ParMETIS_V3_RefineKway");
1880       PetscCall(PetscLogEventBegin(DMPLEX_RebalPartition, 0, 0, 0, 0));
1881       ierr = ParMETIS_V3_RefineKway((PetscInt *)cumSumVertices, (idx_t *)xadj, (idx_t *)adjncy, vtxwgt, NULL, &wgtflag, &numflag, &ncon, &nparts, tpwgts, ubvec, options, &edgecut, part, &comm);
1882       PetscCall(PetscLogEventEnd(DMPLEX_RebalPartition, 0, 0, 0, 0));
1883       PetscStackPop;
1884       PetscCheck(ierr == METIS_OK, PETSC_COMM_SELF, PETSC_ERR_LIB, "Error in ParMETIS_V3_RefineKway()");
1885     } else {
1886       PetscStackPushExternal("ParMETIS_V3_PartKway");
1887       PetscCall(PetscLogEventBegin(DMPLEX_RebalPartition, 0, 0, 0, 0));
1888       ierr = ParMETIS_V3_PartKway((PetscInt *)cumSumVertices, (idx_t *)xadj, (idx_t *)adjncy, vtxwgt, NULL, &wgtflag, &numflag, &ncon, &nparts, tpwgts, ubvec, options, &edgecut, part, &comm);
1889       PetscCall(PetscLogEventEnd(DMPLEX_RebalPartition, 0, 0, 0, 0));
1890       PetscStackPop;
1891       PetscCheck(ierr == METIS_OK, PETSC_COMM_SELF, PETSC_ERR_LIB, "Error in ParMETIS_V3_PartKway()");
1892     }
1893     PetscCall(MatRestoreRowIJ(A, PETSC_FALSE, PETSC_FALSE, PETSC_FALSE, &numRows, &xadj, &adjncy, &done));
1894     PetscCall(PetscFree(options));
1895   } else {
1896     if (viewer) PetscCall(PetscViewerASCIIPrintf(viewer, "Using METIS to partition graph.\n"));
1897     Mat       As;
1898     PetscInt *partGlobal;
1899     PetscInt *numExclusivelyOwnedAll;
1900 
1901     PetscCall(PetscMalloc1(cumSumVertices[rank + 1] - cumSumVertices[rank], &part));
1902     PetscCall(MatGetSize(A, &numRows, NULL));
1903     PetscCall(PetscLogEventBegin(DMPLEX_RebalGatherGraph, dm, 0, 0, 0));
1904     PetscCall(MatMPIAdjToSeqRankZero(A, &As));
1905     PetscCall(PetscLogEventEnd(DMPLEX_RebalGatherGraph, dm, 0, 0, 0));
1906 
1907     PetscCall(PetscMalloc1(size, &numExclusivelyOwnedAll));
1908     numExclusivelyOwnedAll[rank] = numExclusivelyOwned;
1909     PetscCallMPI(MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, numExclusivelyOwnedAll, 1, MPIU_INT, comm));
1910 
1911     PetscCall(PetscMalloc1(numRows, &partGlobal));
1912     PetscCall(PetscLogEventBegin(DMPLEX_RebalPartition, 0, 0, 0, 0));
1913     if (rank == 0) {
1914       PetscInt *vtxwgt_g, numRows_g;
1915 
1916       PetscCall(MatGetRowIJ(As, PETSC_FALSE, PETSC_FALSE, PETSC_FALSE, &numRows_g, &xadj, &adjncy, &done));
1917       PetscCall(PetscMalloc1(2 * numRows_g, &vtxwgt_g));
1918       for (i = 0; i < size; i++) {
1919         vtxwgt_g[ncon * cumSumVertices[i]] = numExclusivelyOwnedAll[i];
1920         if (ncon > 1) vtxwgt_g[ncon * cumSumVertices[i] + 1] = 1;
1921         for (j = cumSumVertices[i] + 1; j < cumSumVertices[i + 1]; j++) {
1922           vtxwgt_g[ncon * j] = 1;
1923           if (ncon > 1) vtxwgt_g[2 * j + 1] = 0;
1924         }
1925       }
1926 
1927       PetscCall(PetscMalloc1(64, &options));
1928       ierr = METIS_SetDefaultOptions(options); /* initialize all defaults */
1929       PetscCheck(ierr == METIS_OK, PETSC_COMM_SELF, PETSC_ERR_LIB, "Error in METIS_SetDefaultOptions()");
1930       options[METIS_OPTION_CONTIG] = 1;
1931       PetscStackPushExternal("METIS_PartGraphKway");
1932       ierr = METIS_PartGraphKway(&numRows_g, &ncon, (idx_t *)xadj, (idx_t *)adjncy, vtxwgt_g, NULL, NULL, &nparts, tpwgts, ubvec, options, &edgecut, partGlobal);
1933       PetscStackPop;
1934       PetscCheck(ierr == METIS_OK, PETSC_COMM_SELF, PETSC_ERR_LIB, "Error in METIS_PartGraphKway()");
1935       PetscCall(PetscFree(options));
1936       PetscCall(PetscFree(vtxwgt_g));
1937       PetscCall(MatRestoreRowIJ(As, PETSC_FALSE, PETSC_FALSE, PETSC_FALSE, &numRows_g, &xadj, &adjncy, &done));
1938       PetscCall(MatDestroy(&As));
1939     }
1940     PetscCall(PetscBarrier((PetscObject)dm));
1941     PetscCall(PetscLogEventEnd(DMPLEX_RebalPartition, 0, 0, 0, 0));
1942     PetscCall(PetscFree(numExclusivelyOwnedAll));
1943 
1944     /* scatter the partitioning information to ranks */
1945     PetscCall(PetscLogEventBegin(DMPLEX_RebalScatterPart, 0, 0, 0, 0));
1946     PetscCall(PetscMalloc1(size, &counts));
1947     PetscCall(PetscMalloc1(size + 1, &mpiCumSumVertices));
1948     for (i = 0; i < size; i++) PetscCall(PetscMPIIntCast(cumSumVertices[i + 1] - cumSumVertices[i], &counts[i]));
1949     for (i = 0; i <= size; i++) PetscCall(PetscMPIIntCast(cumSumVertices[i], &mpiCumSumVertices[i]));
1950     PetscCallMPI(MPI_Scatterv(partGlobal, counts, mpiCumSumVertices, MPIU_INT, part, counts[rank], MPIU_INT, 0, comm));
1951     PetscCall(PetscFree(counts));
1952     PetscCall(PetscFree(mpiCumSumVertices));
1953     PetscCall(PetscFree(partGlobal));
1954     PetscCall(PetscLogEventEnd(DMPLEX_RebalScatterPart, 0, 0, 0, 0));
1955   }
1956   PetscCall(PetscFree(ubvec));
1957   PetscCall(PetscFree(tpwgts));
1958 
1959   /* Rename the result so that the vertex resembling the exclusively owned points stays on the same rank */
1960   PetscCall(PetscMalloc2(size, &firstVertices, size, &renumbering));
1961   firstVertices[rank] = part[0];
1962   PetscCallMPI(MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, firstVertices, 1, MPIU_INT, comm));
1963   for (i = 0; i < size; i++) renumbering[firstVertices[i]] = i;
1964   for (i = 0; i < cumSumVertices[rank + 1] - cumSumVertices[rank]; i++) part[i] = renumbering[part[i]];
1965   PetscCall(PetscFree2(firstVertices, renumbering));
1966 
1967   /* Check if the renumbering worked (this can fail when ParMETIS gives fewer partitions than there are processes) */
1968   failed = (PetscInt)(part[0] != rank);
1969   PetscCall(MPIU_Allreduce(&failed, &failedGlobal, 1, MPIU_INT, MPI_SUM, comm));
1970   if (failedGlobal > 0) {
1971     PetscCheck(failedGlobal <= 0, comm, PETSC_ERR_LIB, "Metis/Parmetis returned a bad partition");
1972     PetscCall(PetscFree(vtxwgt));
1973     PetscCall(PetscFree(toBalance));
1974     PetscCall(PetscFree(isLeaf));
1975     PetscCall(PetscFree(isNonExclusivelyOwned));
1976     PetscCall(PetscFree(isExclusivelyOwned));
1977     if (usematpartitioning) {
1978       PetscCall(ISRestoreIndices(ispart, (const PetscInt **)&part));
1979       PetscCall(ISDestroy(&ispart));
1980     } else PetscCall(PetscFree(part));
1981     if (viewer) {
1982       PetscCall(PetscViewerPopFormat(viewer));
1983       PetscCall(PetscOptionsRestoreViewer(&viewer));
1984     }
1985     PetscCall(PetscLogEventEnd(DMPLEX_RebalanceSharedPoints, dm, 0, 0, 0));
1986     PetscFunctionReturn(PETSC_SUCCESS);
1987   }
1988 
1989   /* Check how well we did distributing points*/
1990   if (viewer) {
1991     PetscCall(PetscViewerASCIIPrintf(viewer, "Number of owned entities of depth %" PetscInt_FMT ".\n", entityDepth));
1992     PetscCall(PetscViewerASCIIPrintf(viewer, "Initial      "));
1993     PetscCall(DMPlexViewDistribution(comm, cumSumVertices[rank + 1] - cumSumVertices[rank], ncon, vtxwgt, NULL, viewer));
1994     PetscCall(PetscViewerASCIIPrintf(viewer, "Rebalanced   "));
1995     PetscCall(DMPlexViewDistribution(comm, cumSumVertices[rank + 1] - cumSumVertices[rank], ncon, vtxwgt, part, viewer));
1996   }
1997 
1998   /* Check that every vertex is owned by a process that it is actually connected to. */
1999   PetscCall(MatGetRowIJ(A, PETSC_FALSE, PETSC_FALSE, PETSC_FALSE, &numRows, (const PetscInt **)&xadj, (const PetscInt **)&adjncy, &done));
2000   for (i = 1; i <= numNonExclusivelyOwned; i++) {
2001     PetscInt loc = 0;
2002     PetscCall(PetscFindInt(cumSumVertices[part[i]], xadj[i + 1] - xadj[i], &adjncy[xadj[i]], &loc));
2003     /* If not, then just set the owner to the original owner (hopefully a rare event, it means that a vertex has been isolated) */
2004     if (loc < 0) part[i] = rank;
2005   }
2006   PetscCall(MatRestoreRowIJ(A, PETSC_FALSE, PETSC_FALSE, PETSC_FALSE, &numRows, (const PetscInt **)&xadj, (const PetscInt **)&adjncy, &done));
2007   PetscCall(MatDestroy(&A));
2008 
2009   /* See how significant the influences of the previous fixing up step was.*/
2010   if (viewer) {
2011     PetscCall(PetscViewerASCIIPrintf(viewer, "After fix    "));
2012     PetscCall(DMPlexViewDistribution(comm, cumSumVertices[rank + 1] - cumSumVertices[rank], ncon, vtxwgt, part, viewer));
2013   }
2014   if (!usematpartitioning) PetscCall(PetscFree(vtxwgt));
2015   else PetscCall(MatPartitioningDestroy(&mp));
2016 
2017   PetscCall(PetscLayoutDestroy(&layout));
2018 
2019   PetscCall(PetscLogEventBegin(DMPLEX_RebalRewriteSF, dm, 0, 0, 0));
2020   /* Rewrite the SF to reflect the new ownership. */
2021   PetscCall(PetscMalloc1(numNonExclusivelyOwned, &pointsToRewrite));
2022   counter = 0;
2023   for (i = 0; i < pEnd - pStart; i++) {
2024     if (toBalance[i]) {
2025       if (isNonExclusivelyOwned[i]) {
2026         pointsToRewrite[counter] = i + pStart;
2027         counter++;
2028       }
2029     }
2030   }
2031   PetscCall(DMPlexRewriteSF(dm, numNonExclusivelyOwned, pointsToRewrite, part + 1, degrees));
2032   PetscCall(PetscFree(pointsToRewrite));
2033   PetscCall(PetscLogEventEnd(DMPLEX_RebalRewriteSF, dm, 0, 0, 0));
2034 
2035   PetscCall(PetscFree(toBalance));
2036   PetscCall(PetscFree(isLeaf));
2037   PetscCall(PetscFree(isNonExclusivelyOwned));
2038   PetscCall(PetscFree(isExclusivelyOwned));
2039   if (usematpartitioning) {
2040     PetscCall(ISRestoreIndices(ispart, (const PetscInt **)&part));
2041     PetscCall(ISDestroy(&ispart));
2042   } else PetscCall(PetscFree(part));
2043   if (viewer) {
2044     PetscCall(PetscViewerPopFormat(viewer));
2045     PetscCall(PetscViewerDestroy(&viewer));
2046   }
2047   if (success) *success = PETSC_TRUE;
2048   PetscCall(PetscLogEventEnd(DMPLEX_RebalanceSharedPoints, dm, 0, 0, 0));
2049   PetscFunctionReturn(PETSC_SUCCESS);
2050 #else
2051   SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Mesh partitioning needs external package support.\nPlease reconfigure with --download-parmetis.");
2052 #endif
2053 }
2054