Actual source code: partmatpart.c

  1: #include <petscmat.h>
  2: #include <petsc/private/partitionerimpl.h>

  4: typedef struct {
  5:   MatPartitioning mp;
  6: } PetscPartitioner_MatPartitioning;

  8: static PetscErrorCode PetscPartitionerMatPartitioningGetMatPartitioning_MatPartitioning(PetscPartitioner part, MatPartitioning *mp)
  9: {
 10:   PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;

 12:   PetscFunctionBegin;
 13:   *mp = p->mp;
 14:   PetscFunctionReturn(PETSC_SUCCESS);
 15: }

 17: /*@
 18:   PetscPartitionerMatPartitioningGetMatPartitioning - Get a `MatPartitioning` instance wrapped by this `PetscPartitioner`.

 20:   Not Collective

 22:   Input Parameter:
 23: . part - The `PetscPartitioner`

 25:   Output Parameter:
 26: . mp - The `MatPartitioning`

 28:   Level: developer

 30: .seealso: `DMPlexDistribute()`, `PetscPartitionerCreate()`
 31: @*/
 32: PetscErrorCode PetscPartitionerMatPartitioningGetMatPartitioning(PetscPartitioner part, MatPartitioning *mp)
 33: {
 34:   PetscFunctionBegin;
 36:   PetscAssertPointer(mp, 2);
 37:   PetscUseMethod(part, "PetscPartitionerMatPartitioningGetMatPartitioning_C", (PetscPartitioner, MatPartitioning *), (part, mp));
 38:   PetscFunctionReturn(PETSC_SUCCESS);
 39: }

 41: static PetscErrorCode PetscPartitionerDestroy_MatPartitioning(PetscPartitioner part)
 42: {
 43:   PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;

 45:   PetscFunctionBegin;
 46:   PetscCall(MatPartitioningDestroy(&p->mp));
 47:   PetscCall(PetscObjectComposeFunction((PetscObject)part, "PetscPartitionerMatPartitioningGetMatPartitioning_C", NULL));
 48:   PetscCall(PetscFree(part->data));
 49:   PetscFunctionReturn(PETSC_SUCCESS);
 50: }

 52: static PetscErrorCode PetscPartitionerView_MatPartitioning_ASCII(PetscPartitioner part, PetscViewer viewer)
 53: {
 54:   PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;
 55:   PetscViewerFormat                 format;

 57:   PetscFunctionBegin;
 58:   PetscCall(PetscViewerGetFormat(viewer, &format));
 59:   PetscCall(PetscViewerASCIIPrintf(viewer, "MatPartitioning Graph Partitioner:\n"));
 60:   PetscCall(PetscViewerASCIIPushTab(viewer));
 61:   if (p->mp) PetscCall(MatPartitioningView(p->mp, viewer));
 62:   PetscCall(PetscViewerASCIIPopTab(viewer));
 63:   PetscFunctionReturn(PETSC_SUCCESS);
 64: }

 66: static PetscErrorCode PetscPartitionerView_MatPartitioning(PetscPartitioner part, PetscViewer viewer)
 67: {
 68:   PetscBool iascii;

 70:   PetscFunctionBegin;
 73:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
 74:   if (iascii) PetscCall(PetscPartitionerView_MatPartitioning_ASCII(part, viewer));
 75:   PetscFunctionReturn(PETSC_SUCCESS);
 76: }

 78: static PetscErrorCode PetscPartitionerSetFromOptions_MatPartitioning(PetscPartitioner part, PetscOptionItems *PetscOptionsObject)
 79: {
 80:   PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;

 82:   PetscFunctionBegin;
 83:   PetscCall(PetscObjectSetOptionsPrefix((PetscObject)p->mp, ((PetscObject)part)->prefix));
 84:   PetscCall(MatPartitioningSetFromOptions(p->mp));
 85:   PetscFunctionReturn(PETSC_SUCCESS);
 86: }

 88: static PetscErrorCode PetscPartitionerPartition_MatPartitioning(PetscPartitioner part, PetscInt nparts, PetscInt numVertices, PetscInt start[], PetscInt adjacency[], PetscSection vertSection, PetscSection edgeSection, PetscSection targetSection, PetscSection partSection, IS *is)
 89: {
 90:   PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;
 91:   Mat                               matadj;
 92:   IS                                is1, is2, is3;
 93:   PetscReal                        *tpwgts = NULL;
 94:   PetscInt                          numVerticesGlobal, numEdges;
 95:   PetscInt                         *i, *j, *vwgt = NULL;
 96:   MPI_Comm                          comm;

 98:   PetscFunctionBegin;
 99:   PetscCall(PetscObjectGetComm((PetscObject)part, &comm));

101:   /* TODO: MatCreateMPIAdj should maybe take global number of ROWS */
102:   /* TODO: And vertex distribution in PetscPartitionerPartition_ParMetis should be done using PetscSplitOwnership */
103:   numVerticesGlobal = PETSC_DECIDE;
104:   PetscCall(PetscSplitOwnership(comm, &numVertices, &numVerticesGlobal));

106:   /* copy arrays to avoid memory errors because MatMPIAdjSetPreallocation copies just pointers */
107:   numEdges = start[numVertices];
108:   PetscCall(PetscMalloc1(numVertices + 1, &i));
109:   PetscCall(PetscMalloc1(numEdges, &j));
110:   PetscCall(PetscArraycpy(i, start, numVertices + 1));
111:   PetscCall(PetscArraycpy(j, adjacency, numEdges));

113:   /* construct the adjacency matrix */
114:   PetscCall(MatCreateMPIAdj(comm, numVertices, numVerticesGlobal, i, j, NULL, &matadj));
115:   PetscCall(MatPartitioningSetAdjacency(p->mp, matadj));
116:   PetscCall(MatPartitioningSetNParts(p->mp, nparts));

118:   /* calculate partition weights */
119:   if (targetSection) {
120:     PetscReal sumt;
121:     PetscInt  p;

123:     sumt = 0.0;
124:     PetscCall(PetscMalloc1(nparts, &tpwgts));
125:     for (p = 0; p < nparts; ++p) {
126:       PetscInt tpd;

128:       PetscCall(PetscSectionGetDof(targetSection, p, &tpd));
129:       sumt += tpd;
130:       tpwgts[p] = tpd;
131:     }
132:     if (sumt) { /* METIS/ParMETIS do not like exactly zero weight */
133:       for (p = 0, sumt = 0.0; p < nparts; ++p) {
134:         tpwgts[p] = PetscMax(tpwgts[p], PETSC_SMALL);
135:         sumt += tpwgts[p];
136:       }
137:       for (p = 0; p < nparts; ++p) tpwgts[p] /= sumt;
138:       for (p = 0, sumt = 0.0; p < nparts - 1; ++p) sumt += tpwgts[p];
139:       tpwgts[nparts - 1] = 1. - sumt;
140:     } else {
141:       PetscCall(PetscFree(tpwgts));
142:     }
143:   }
144:   PetscCall(MatPartitioningSetPartitionWeights(p->mp, tpwgts));

146:   /* calculate vertex weights */
147:   if (vertSection) {
148:     PetscInt v;

150:     PetscCall(PetscMalloc1(numVertices, &vwgt));
151:     for (v = 0; v < numVertices; ++v) PetscCall(PetscSectionGetDof(vertSection, v, &vwgt[v]));
152:   }
153:   PetscCall(MatPartitioningSetVertexWeights(p->mp, vwgt));

155:   /* apply the partitioning */
156:   PetscCall(MatPartitioningApply(p->mp, &is1));

158:   /* construct the PetscSection */
159:   {
160:     PetscInt        v;
161:     const PetscInt *assignment_arr;

163:     PetscCall(ISGetIndices(is1, &assignment_arr));
164:     for (v = 0; v < numVertices; ++v) PetscCall(PetscSectionAddDof(partSection, assignment_arr[v], 1));
165:     PetscCall(ISRestoreIndices(is1, &assignment_arr));
166:   }

168:   /* convert assignment IS to global numbering IS */
169:   PetscCall(ISPartitioningToNumbering(is1, &is2));
170:   PetscCall(ISDestroy(&is1));

172:   /* renumber IS into local numbering */
173:   PetscCall(ISOnComm(is2, PETSC_COMM_SELF, PETSC_USE_POINTER, &is1));
174:   PetscCall(ISRenumber(is1, NULL, NULL, &is3));
175:   PetscCall(ISDestroy(&is1));
176:   PetscCall(ISDestroy(&is2));

178:   /* invert IS */
179:   PetscCall(ISSetPermutation(is3));
180:   PetscCall(ISInvertPermutation(is3, numVertices, &is1));
181:   PetscCall(ISDestroy(&is3));

183:   PetscCall(MatDestroy(&matadj));
184:   *is = is1;
185:   PetscFunctionReturn(PETSC_SUCCESS);
186: }

188: static PetscErrorCode PetscPartitionerInitialize_MatPartitioning(PetscPartitioner part)
189: {
190:   PetscFunctionBegin;
191:   part->ops->view           = PetscPartitionerView_MatPartitioning;
192:   part->ops->setfromoptions = PetscPartitionerSetFromOptions_MatPartitioning;
193:   part->ops->destroy        = PetscPartitionerDestroy_MatPartitioning;
194:   part->ops->partition      = PetscPartitionerPartition_MatPartitioning;
195:   PetscCall(PetscObjectComposeFunction((PetscObject)part, "PetscPartitionerMatPartitioningGetMatPartitioning_C", PetscPartitionerMatPartitioningGetMatPartitioning_MatPartitioning));
196:   PetscFunctionReturn(PETSC_SUCCESS);
197: }

199: /*MC
200:   PETSCPARTITIONERMATPARTITIONING = "matpartitioning" - A PetscPartitioner object

202:   Level: developer

204: .seealso: `PetscPartitionerType`, `PetscPartitionerCreate()`, `PetscPartitionerSetType()`
205: M*/

207: PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_MatPartitioning(PetscPartitioner part)
208: {
209:   PetscPartitioner_MatPartitioning *p;

211:   PetscFunctionBegin;
213:   PetscCall(PetscNew(&p));
214:   part->data = p;
215:   PetscCall(PetscPartitionerInitialize_MatPartitioning(part));
216:   PetscCall(MatPartitioningCreate(PetscObjectComm((PetscObject)part), &p->mp));
217:   PetscFunctionReturn(PETSC_SUCCESS);
218: }