1 #include <petscmat.h>
2 #include <petsc/private/partitionerimpl.h> /*I "petscpartitioner.h" I*/
3
4 typedef struct {
5 MatPartitioning mp;
6 } PetscPartitioner_MatPartitioning;
7
PetscPartitionerMatPartitioningGetMatPartitioning_MatPartitioning(PetscPartitioner part,MatPartitioning * mp)8 static PetscErrorCode PetscPartitionerMatPartitioningGetMatPartitioning_MatPartitioning(PetscPartitioner part, MatPartitioning *mp)
9 {
10 PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;
11
12 PetscFunctionBegin;
13 *mp = p->mp;
14 PetscFunctionReturn(PETSC_SUCCESS);
15 }
16
17 /*@
18 PetscPartitionerMatPartitioningGetMatPartitioning - Get a `MatPartitioning` instance wrapped by this `PetscPartitioner`.
19
20 Not Collective
21
22 Input Parameter:
23 . part - The `PetscPartitioner`
24
25 Output Parameter:
26 . mp - The `MatPartitioning`
27
28 Level: developer
29
30 .seealso: `DMPlexDistribute()`, `PetscPartitionerCreate()`
31 @*/
PetscPartitionerMatPartitioningGetMatPartitioning(PetscPartitioner part,MatPartitioning * mp)32 PetscErrorCode PetscPartitionerMatPartitioningGetMatPartitioning(PetscPartitioner part, MatPartitioning *mp)
33 {
34 PetscFunctionBegin;
35 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1);
36 PetscAssertPointer(mp, 2);
37 PetscUseMethod(part, "PetscPartitionerMatPartitioningGetMatPartitioning_C", (PetscPartitioner, MatPartitioning *), (part, mp));
38 PetscFunctionReturn(PETSC_SUCCESS);
39 }
40
PetscPartitionerDestroy_MatPartitioning(PetscPartitioner part)41 static PetscErrorCode PetscPartitionerDestroy_MatPartitioning(PetscPartitioner part)
42 {
43 PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;
44
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 }
51
PetscPartitionerView_MatPartitioning_ASCII(PetscPartitioner part,PetscViewer viewer)52 static PetscErrorCode PetscPartitionerView_MatPartitioning_ASCII(PetscPartitioner part, PetscViewer viewer)
53 {
54 PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;
55 PetscViewerFormat format;
56
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 }
65
PetscPartitionerView_MatPartitioning(PetscPartitioner part,PetscViewer viewer)66 static PetscErrorCode PetscPartitionerView_MatPartitioning(PetscPartitioner part, PetscViewer viewer)
67 {
68 PetscBool isascii;
69
70 PetscFunctionBegin;
71 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1);
72 PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
73 PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
74 if (isascii) PetscCall(PetscPartitionerView_MatPartitioning_ASCII(part, viewer));
75 PetscFunctionReturn(PETSC_SUCCESS);
76 }
77
PetscPartitionerSetFromOptions_MatPartitioning(PetscPartitioner part,PetscOptionItems PetscOptionsObject)78 static PetscErrorCode PetscPartitionerSetFromOptions_MatPartitioning(PetscPartitioner part, PetscOptionItems PetscOptionsObject)
79 {
80 PetscPartitioner_MatPartitioning *p = (PetscPartitioner_MatPartitioning *)part->data;
81
82 PetscFunctionBegin;
83 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)p->mp, ((PetscObject)part)->prefix));
84 PetscCall(MatPartitioningSetFromOptions(p->mp));
85 PetscFunctionReturn(PETSC_SUCCESS);
86 }
87
PetscPartitionerPartition_MatPartitioning(PetscPartitioner part,PetscInt nparts,PetscInt numVertices,PetscInt start[],PetscInt adjacency[],PetscSection vertSection,PetscSection edgeSection,PetscSection targetSection,PetscSection partSection,IS * is)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;
97
98 PetscFunctionBegin;
99 PetscCall(PetscObjectGetComm((PetscObject)part, &comm));
100
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));
105
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));
112
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));
117
118 /* calculate partition weights */
119 if (targetSection) {
120 PetscReal sumt;
121 PetscInt p;
122
123 sumt = 0.0;
124 PetscCall(PetscMalloc1(nparts, &tpwgts));
125 for (p = 0; p < nparts; ++p) {
126 PetscInt tpd;
127
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));
145
146 /* calculate vertex weights */
147 if (vertSection) {
148 PetscInt v;
149
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));
154
155 /* apply the partitioning */
156 PetscCall(MatPartitioningApply(p->mp, &is1));
157
158 /* construct the PetscSection */
159 {
160 PetscInt v;
161 const PetscInt *assignment_arr;
162
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 }
167
168 /* convert assignment IS to global numbering IS */
169 PetscCall(ISPartitioningToNumbering(is1, &is2));
170 PetscCall(ISDestroy(&is1));
171
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));
177
178 /* invert IS */
179 PetscCall(ISSetPermutation(is3));
180 PetscCall(ISInvertPermutation(is3, numVertices, &is1));
181 PetscCall(ISDestroy(&is3));
182
183 PetscCall(MatDestroy(&matadj));
184 *is = is1;
185 PetscFunctionReturn(PETSC_SUCCESS);
186 }
187
PetscPartitionerInitialize_MatPartitioning(PetscPartitioner part)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 }
198
199 /*MC
200 PETSCPARTITIONERMATPARTITIONING = "matpartitioning" - A PetscPartitioner object
201
202 Level: developer
203
204 .seealso: `PetscPartitionerType`, `PetscPartitionerCreate()`, `PetscPartitionerSetType()`
205 M*/
206
PetscPartitionerCreate_MatPartitioning(PetscPartitioner part)207 PETSC_EXTERN PetscErrorCode PetscPartitionerCreate_MatPartitioning(PetscPartitioner part)
208 {
209 PetscPartitioner_MatPartitioning *p;
210
211 PetscFunctionBegin;
212 PetscValidHeaderSpecific(part, PETSCPARTITIONER_CLASSID, 1);
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 }
219