1 #include <../src/mat/impls/aij/mpi/mpiaij.h> /*I "petscmat.h" I*/
2 #include <petscsf.h>
3
4 typedef struct {
5 PetscSF sf;
6 PetscReal *dwts, *owts;
7 PetscInt *dmask, *omask, *cmask;
8 PetscBool local;
9 } MC_JP;
10
MatColoringDestroy_JP(MatColoring mc)11 static PetscErrorCode MatColoringDestroy_JP(MatColoring mc)
12 {
13 PetscFunctionBegin;
14 PetscCall(PetscFree(mc->data));
15 PetscFunctionReturn(PETSC_SUCCESS);
16 }
17
MatColoringSetFromOptions_JP(MatColoring mc,PetscOptionItems PetscOptionsObject)18 static PetscErrorCode MatColoringSetFromOptions_JP(MatColoring mc, PetscOptionItems PetscOptionsObject)
19 {
20 MC_JP *jp = (MC_JP *)mc->data;
21
22 PetscFunctionBegin;
23 PetscOptionsHeadBegin(PetscOptionsObject, "JP options");
24 PetscCall(PetscOptionsBool("-mat_coloring_jp_local", "Do an initial coloring of local columns", "", jp->local, &jp->local, NULL));
25 PetscOptionsHeadEnd();
26 PetscFunctionReturn(PETSC_SUCCESS);
27 }
28
MCJPGreatestWeight_Private(MatColoring mc,const PetscReal * weights,PetscReal * maxweights)29 static PetscErrorCode MCJPGreatestWeight_Private(MatColoring mc, const PetscReal *weights, PetscReal *maxweights)
30 {
31 MC_JP *jp = (MC_JP *)mc->data;
32 Mat G = mc->mat, dG, oG;
33 PetscBool isSeq, isMPI;
34 Mat_MPIAIJ *aij;
35 Mat_SeqAIJ *daij, *oaij;
36 PetscInt *di, *oi, *dj, *oj;
37 PetscSF sf = jp->sf;
38 PetscLayout layout;
39 PetscInt dn, on;
40 PetscInt i, j, l;
41 PetscReal *dwts = jp->dwts, *owts = jp->owts;
42 PetscInt ncols;
43 const PetscInt *cols;
44
45 PetscFunctionBegin;
46 PetscCall(PetscObjectTypeCompare((PetscObject)G, MATSEQAIJ, &isSeq));
47 PetscCall(PetscObjectTypeCompare((PetscObject)G, MATMPIAIJ, &isMPI));
48 PetscCheck(isSeq || isMPI, PetscObjectComm((PetscObject)G), PETSC_ERR_ARG_WRONGSTATE, "MatColoringDegrees requires an MPI/SEQAIJ Matrix");
49
50 /* get the inner matrix nonzero structure */
51 oG = NULL;
52 oi = NULL;
53 oj = NULL;
54 if (isMPI) {
55 aij = (Mat_MPIAIJ *)G->data;
56 dG = aij->A;
57 oG = aij->B;
58 daij = (Mat_SeqAIJ *)dG->data;
59 oaij = (Mat_SeqAIJ *)oG->data;
60 di = daij->i;
61 dj = daij->j;
62 oi = oaij->i;
63 oj = oaij->j;
64 PetscCall(MatGetSize(oG, &dn, &on));
65 if (!sf) {
66 PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)mc), &sf));
67 PetscCall(MatGetLayouts(G, &layout, NULL));
68 PetscCall(PetscSFSetGraphLayout(sf, layout, on, NULL, PETSC_COPY_VALUES, aij->garray));
69 jp->sf = sf;
70 }
71 } else {
72 dG = G;
73 PetscCall(MatGetSize(dG, NULL, &dn));
74 daij = (Mat_SeqAIJ *)dG->data;
75 di = daij->i;
76 dj = daij->j;
77 }
78 /* set up the distance-zero weights */
79 if (!dwts) {
80 PetscCall(PetscMalloc1(dn, &dwts));
81 jp->dwts = dwts;
82 if (oG) {
83 PetscCall(PetscMalloc1(on, &owts));
84 jp->owts = owts;
85 }
86 }
87 for (i = 0; i < dn; i++) {
88 maxweights[i] = weights[i];
89 dwts[i] = maxweights[i];
90 }
91 /* get the off-diagonal weights */
92 if (oG) {
93 PetscCall(PetscLogEventBegin(MATCOLORING_Comm, mc, 0, 0, 0));
94 PetscCall(PetscSFBcastBegin(sf, MPIU_REAL, dwts, owts, MPI_REPLACE));
95 PetscCall(PetscSFBcastEnd(sf, MPIU_REAL, dwts, owts, MPI_REPLACE));
96 PetscCall(PetscLogEventEnd(MATCOLORING_Comm, mc, 0, 0, 0));
97 }
98 /* check for the maximum out to the distance of the coloring */
99 for (l = 0; l < mc->dist; l++) {
100 /* check for on-diagonal greater weights */
101 for (i = 0; i < dn; i++) {
102 ncols = di[i + 1] - di[i];
103 cols = &(dj[di[i]]);
104 for (j = 0; j < ncols; j++) {
105 if (dwts[cols[j]] > maxweights[i]) maxweights[i] = dwts[cols[j]];
106 }
107 /* check for off-diagonal greater weights */
108 if (oG) {
109 ncols = oi[i + 1] - oi[i];
110 cols = &(oj[oi[i]]);
111 for (j = 0; j < ncols; j++) {
112 if (owts[cols[j]] > maxweights[i]) maxweights[i] = owts[cols[j]];
113 }
114 }
115 }
116 if (l < mc->dist - 1) {
117 for (i = 0; i < dn; i++) dwts[i] = maxweights[i];
118 if (oG) {
119 PetscCall(PetscLogEventBegin(MATCOLORING_Comm, mc, 0, 0, 0));
120 PetscCall(PetscSFBcastBegin(sf, MPIU_REAL, dwts, owts, MPI_REPLACE));
121 PetscCall(PetscSFBcastEnd(sf, MPIU_REAL, dwts, owts, MPI_REPLACE));
122 PetscCall(PetscLogEventEnd(MATCOLORING_Comm, mc, 0, 0, 0));
123 }
124 }
125 }
126 PetscFunctionReturn(PETSC_SUCCESS);
127 }
128
MCJPInitialLocalColor_Private(MatColoring mc,PetscInt * lperm,ISColoringValue * colors)129 static PetscErrorCode MCJPInitialLocalColor_Private(MatColoring mc, PetscInt *lperm, ISColoringValue *colors)
130 {
131 PetscInt j, i, s, e, n, bidx, cidx, idx, dist, distance = mc->dist;
132 Mat G = mc->mat, dG, oG;
133 PetscInt *seen;
134 PetscInt *idxbuf;
135 PetscBool *boundary;
136 PetscInt *distbuf;
137 PetscInt *colormask;
138 PetscInt ncols;
139 const PetscInt *cols;
140 PetscBool isSeq, isMPI;
141 Mat_MPIAIJ *aij;
142 Mat_SeqAIJ *daij, *oaij;
143 PetscInt *di, *dj, dn;
144 PetscInt *oi;
145
146 PetscFunctionBegin;
147 PetscCall(PetscLogEventBegin(MATCOLORING_Local, mc, 0, 0, 0));
148 PetscCall(MatGetOwnershipRange(G, &s, &e));
149 n = e - s;
150 PetscCall(PetscObjectBaseTypeCompare((PetscObject)G, MATSEQAIJ, &isSeq));
151 PetscCall(PetscObjectTypeCompare((PetscObject)G, MATMPIAIJ, &isMPI));
152 PetscCheck(isSeq || isMPI, PetscObjectComm((PetscObject)G), PETSC_ERR_ARG_WRONGSTATE, "MatColoringDegrees requires an MPI/SEQAIJ Matrix");
153
154 /* get the inner matrix nonzero structure */
155 oG = NULL;
156 oi = NULL;
157 if (isMPI) {
158 aij = (Mat_MPIAIJ *)G->data;
159 dG = aij->A;
160 oG = aij->B;
161 daij = (Mat_SeqAIJ *)dG->data;
162 oaij = (Mat_SeqAIJ *)oG->data;
163 di = daij->i;
164 dj = daij->j;
165 oi = oaij->i;
166 PetscCall(MatGetSize(oG, &dn, NULL));
167 } else {
168 dG = G;
169 PetscCall(MatGetSize(dG, NULL, &dn));
170 daij = (Mat_SeqAIJ *)dG->data;
171 di = daij->i;
172 dj = daij->j;
173 }
174 PetscCall(PetscMalloc5(n, &colormask, n, &seen, n, &idxbuf, n, &distbuf, n, &boundary));
175 for (i = 0; i < dn; i++) {
176 seen[i] = -1;
177 colormask[i] = -1;
178 boundary[i] = PETSC_FALSE;
179 }
180 /* pass one -- figure out which ones are off-boundary in the distance-n sense */
181 if (oG) {
182 for (i = 0; i < dn; i++) {
183 bidx = -1;
184 /* nonempty off-diagonal, so this one is on the boundary */
185 if (oi[i] != oi[i + 1]) {
186 boundary[i] = PETSC_TRUE;
187 continue;
188 }
189 ncols = di[i + 1] - di[i];
190 cols = &(dj[di[i]]);
191 for (j = 0; j < ncols; j++) {
192 bidx++;
193 seen[cols[j]] = i;
194 distbuf[bidx] = 1;
195 idxbuf[bidx] = cols[j];
196 }
197 while (bidx >= 0) {
198 idx = idxbuf[bidx];
199 dist = distbuf[bidx];
200 bidx--;
201 if (dist < distance) {
202 if (oi[idx + 1] != oi[idx]) {
203 boundary[i] = PETSC_TRUE;
204 break;
205 }
206 ncols = di[idx + 1] - di[idx];
207 cols = &(dj[di[idx]]);
208 for (j = 0; j < ncols; j++) {
209 if (seen[cols[j]] != i) {
210 bidx++;
211 seen[cols[j]] = i;
212 idxbuf[bidx] = cols[j];
213 distbuf[bidx] = dist + 1;
214 }
215 }
216 }
217 }
218 }
219 for (i = 0; i < dn; i++) seen[i] = -1;
220 }
221 /* pass two -- color it by looking at nearby vertices and building a mask */
222 for (i = 0; i < dn; i++) {
223 cidx = lperm[i];
224 if (!boundary[cidx]) {
225 bidx = -1;
226 ncols = di[cidx + 1] - di[cidx];
227 cols = &(dj[di[cidx]]);
228 for (j = 0; j < ncols; j++) {
229 bidx++;
230 seen[cols[j]] = cidx;
231 distbuf[bidx] = 1;
232 idxbuf[bidx] = cols[j];
233 }
234 while (bidx >= 0) {
235 idx = idxbuf[bidx];
236 dist = distbuf[bidx];
237 bidx--;
238 /* mask this color */
239 if (colors[idx] < IS_COLORING_MAX) colormask[colors[idx]] = cidx;
240 if (dist < distance) {
241 ncols = di[idx + 1] - di[idx];
242 cols = &(dj[di[idx]]);
243 for (j = 0; j < ncols; j++) {
244 if (seen[cols[j]] != cidx) {
245 bidx++;
246 seen[cols[j]] = cidx;
247 idxbuf[bidx] = cols[j];
248 distbuf[bidx] = dist + 1;
249 }
250 }
251 }
252 }
253 /* find the lowest untaken color */
254 for (j = 0; j < n; j++) {
255 if (colormask[j] != cidx || j >= mc->maxcolors) {
256 PetscCall(ISColoringValueCast(j, &colors[cidx]));
257 break;
258 }
259 }
260 }
261 }
262 PetscCall(PetscFree5(colormask, seen, idxbuf, distbuf, boundary));
263 PetscCall(PetscLogEventEnd(MATCOLORING_Local, mc, 0, 0, 0));
264 PetscFunctionReturn(PETSC_SUCCESS);
265 }
266
MCJPMinColor_Private(MatColoring mc,ISColoringValue maxcolor,const ISColoringValue * colors,ISColoringValue * mincolors)267 static PetscErrorCode MCJPMinColor_Private(MatColoring mc, ISColoringValue maxcolor, const ISColoringValue *colors, ISColoringValue *mincolors)
268 {
269 MC_JP *jp = (MC_JP *)mc->data;
270 Mat G = mc->mat, dG, oG;
271 PetscBool isSeq, isMPI;
272 Mat_MPIAIJ *aij;
273 Mat_SeqAIJ *daij, *oaij;
274 PetscInt *di, *oi, *dj, *oj;
275 PetscSF sf = jp->sf;
276 PetscLayout layout;
277 PetscInt maskrounds, maskbase, maskradix;
278 PetscInt dn, on;
279 PetscInt i, j, l, k;
280 PetscInt *dmask = jp->dmask, *omask = jp->omask, *cmask = jp->cmask, curmask;
281 PetscInt ncols;
282 const PetscInt *cols;
283
284 PetscFunctionBegin;
285 maskradix = sizeof(PetscInt) * 8;
286 maskrounds = 1 + maxcolor / (maskradix);
287 maskbase = 0;
288 PetscCall(PetscObjectBaseTypeCompare((PetscObject)G, MATSEQAIJ, &isSeq));
289 PetscCall(PetscObjectTypeCompare((PetscObject)G, MATMPIAIJ, &isMPI));
290 PetscCheck(isSeq || isMPI, PetscObjectComm((PetscObject)G), PETSC_ERR_ARG_WRONGSTATE, "MatColoringDegrees requires an MPI/SEQAIJ Matrix");
291
292 /* get the inner matrix nonzero structure */
293 oG = NULL;
294 oi = NULL;
295 oj = NULL;
296 if (isMPI) {
297 aij = (Mat_MPIAIJ *)G->data;
298 dG = aij->A;
299 oG = aij->B;
300 daij = (Mat_SeqAIJ *)dG->data;
301 oaij = (Mat_SeqAIJ *)oG->data;
302 di = daij->i;
303 dj = daij->j;
304 oi = oaij->i;
305 oj = oaij->j;
306 PetscCall(MatGetSize(oG, &dn, &on));
307 if (!sf) {
308 PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)mc), &sf));
309 PetscCall(MatGetLayouts(G, &layout, NULL));
310 PetscCall(PetscSFSetGraphLayout(sf, layout, on, NULL, PETSC_COPY_VALUES, aij->garray));
311 jp->sf = sf;
312 }
313 } else {
314 dG = G;
315 PetscCall(MatGetSize(dG, NULL, &dn));
316 daij = (Mat_SeqAIJ *)dG->data;
317 di = daij->i;
318 dj = daij->j;
319 }
320 for (i = 0; i < dn; i++) mincolors[i] = IS_COLORING_MAX;
321 /* set up the distance-zero mask */
322 if (!dmask) {
323 PetscCall(PetscMalloc1(dn, &dmask));
324 PetscCall(PetscMalloc1(dn, &cmask));
325 jp->cmask = cmask;
326 jp->dmask = dmask;
327 if (oG) {
328 PetscCall(PetscMalloc1(on, &omask));
329 jp->omask = omask;
330 }
331 }
332 /* the number of colors may be more than the number of bits in a PetscInt; take multiple rounds */
333 for (k = 0; k < maskrounds; k++) {
334 for (i = 0; i < dn; i++) {
335 cmask[i] = 0;
336 if (colors[i] < maskbase + maskradix && colors[i] >= maskbase) cmask[i] = 1 << (colors[i] - maskbase);
337 dmask[i] = cmask[i];
338 }
339 if (oG) {
340 PetscCall(PetscLogEventBegin(MATCOLORING_Comm, mc, 0, 0, 0));
341 PetscCall(PetscSFBcastBegin(sf, MPIU_INT, dmask, omask, MPI_REPLACE));
342 PetscCall(PetscSFBcastEnd(sf, MPIU_INT, dmask, omask, MPI_REPLACE));
343 PetscCall(PetscLogEventEnd(MATCOLORING_Comm, mc, 0, 0, 0));
344 }
345 /* fill in the mask out to the distance of the coloring */
346 for (l = 0; l < mc->dist; l++) {
347 /* fill in the on-and-off diagonal mask */
348 for (i = 0; i < dn; i++) {
349 ncols = di[i + 1] - di[i];
350 cols = &(dj[di[i]]);
351 for (j = 0; j < ncols; j++) cmask[i] = cmask[i] | dmask[cols[j]];
352 if (oG) {
353 ncols = oi[i + 1] - oi[i];
354 cols = &(oj[oi[i]]);
355 for (j = 0; j < ncols; j++) cmask[i] = cmask[i] | omask[cols[j]];
356 }
357 }
358 for (i = 0; i < dn; i++) dmask[i] = cmask[i];
359 if (l < mc->dist - 1) {
360 if (oG) {
361 PetscCall(PetscLogEventBegin(MATCOLORING_Comm, mc, 0, 0, 0));
362 PetscCall(PetscSFBcastBegin(sf, MPIU_INT, dmask, omask, MPI_REPLACE));
363 PetscCall(PetscSFBcastEnd(sf, MPIU_INT, dmask, omask, MPI_REPLACE));
364 PetscCall(PetscLogEventEnd(MATCOLORING_Comm, mc, 0, 0, 0));
365 }
366 }
367 }
368 /* read through the mask to see if we've discovered an acceptable color for any vertices in this round */
369 for (i = 0; i < dn; i++) {
370 if (mincolors[i] == IS_COLORING_MAX) {
371 curmask = dmask[i];
372 for (j = 0; j < maskradix; j++) {
373 if (curmask % 2 == 0) {
374 PetscCall(ISColoringValueCast(j + maskbase, &mincolors[i]));
375 break;
376 }
377 curmask = curmask >> 1;
378 }
379 }
380 }
381 /* do the next maskradix colors */
382 maskbase += maskradix;
383 }
384 for (i = 0; i < dn; i++) {
385 if (mincolors[i] == IS_COLORING_MAX) mincolors[i] = maxcolor + 1;
386 }
387 PetscFunctionReturn(PETSC_SUCCESS);
388 }
389
MatColoringApply_JP(MatColoring mc,ISColoring * iscoloring)390 static PetscErrorCode MatColoringApply_JP(MatColoring mc, ISColoring *iscoloring)
391 {
392 MC_JP *jp = (MC_JP *)mc->data;
393 PetscInt i, nadded, nadded_total, nadded_total_old, ntotal, n;
394 PetscInt maxcolor_local = 0, maxcolor_global = 0, *lperm;
395 PetscMPIInt rank;
396 PetscReal *weights, *maxweights;
397 ISColoringValue *color, *mincolor;
398
399 PetscFunctionBegin;
400 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mc), &rank));
401 PetscCall(PetscLogEventBegin(MATCOLORING_Weights, mc, 0, 0, 0));
402 PetscCall(MatColoringCreateWeights(mc, &weights, &lperm));
403 PetscCall(PetscLogEventEnd(MATCOLORING_Weights, mc, 0, 0, 0));
404 PetscCall(MatGetSize(mc->mat, NULL, &ntotal));
405 PetscCall(MatGetLocalSize(mc->mat, NULL, &n));
406 PetscCall(PetscMalloc1(n, &maxweights));
407 PetscCall(PetscMalloc1(n, &color));
408 PetscCall(PetscMalloc1(n, &mincolor));
409 for (i = 0; i < n; i++) {
410 color[i] = IS_COLORING_MAX;
411 mincolor[i] = 0;
412 }
413 nadded = 0;
414 nadded_total = 0;
415 nadded_total_old = 0;
416 /* compute purely local vertices */
417 if (jp->local) {
418 PetscCall(MCJPInitialLocalColor_Private(mc, lperm, color));
419 for (i = 0; i < n; i++) {
420 if (color[i] < IS_COLORING_MAX) {
421 nadded++;
422 weights[i] = -1;
423 if (color[i] > maxcolor_local) maxcolor_local = color[i];
424 }
425 }
426 PetscCallMPI(MPIU_Allreduce(&nadded, &nadded_total, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)mc)));
427 PetscCallMPI(MPIU_Allreduce(&maxcolor_local, &maxcolor_global, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)mc)));
428 }
429 while (nadded_total < ntotal) {
430 PetscCall(MCJPMinColor_Private(mc, (ISColoringValue)maxcolor_global, color, mincolor));
431 PetscCall(MCJPGreatestWeight_Private(mc, weights, maxweights));
432 for (i = 0; i < n; i++) {
433 /* choose locally maximal vertices; weights less than zero are omitted from the graph */
434 if (weights[i] >= maxweights[i] && weights[i] >= 0.) {
435 /* assign the minimum possible color */
436 if (mc->maxcolors > mincolor[i]) {
437 color[i] = mincolor[i];
438 } else {
439 color[i] = (ISColoringValue)mc->maxcolors;
440 }
441 if (color[i] > maxcolor_local) maxcolor_local = color[i];
442 weights[i] = -1.;
443 nadded++;
444 }
445 }
446 PetscCallMPI(MPIU_Allreduce(&maxcolor_local, &maxcolor_global, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)mc)));
447 PetscCallMPI(MPIU_Allreduce(&nadded, &nadded_total, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)mc)));
448 PetscCheck(nadded_total != nadded_total_old, PetscObjectComm((PetscObject)mc), PETSC_ERR_NOT_CONVERGED, "JP didn't make progress");
449 nadded_total_old = nadded_total;
450 }
451 PetscCall(PetscLogEventBegin(MATCOLORING_ISCreate, mc, 0, 0, 0));
452 PetscCall(ISColoringCreate(PetscObjectComm((PetscObject)mc), maxcolor_global + 1, n, color, PETSC_OWN_POINTER, iscoloring));
453 PetscCall(PetscLogEventEnd(MATCOLORING_ISCreate, mc, 0, 0, 0));
454 PetscCall(PetscFree(jp->dwts));
455 PetscCall(PetscFree(jp->dmask));
456 PetscCall(PetscFree(jp->cmask));
457 PetscCall(PetscFree(jp->owts));
458 PetscCall(PetscFree(jp->omask));
459 PetscCall(PetscFree(weights));
460 PetscCall(PetscFree(lperm));
461 PetscCall(PetscFree(maxweights));
462 PetscCall(PetscFree(mincolor));
463 PetscCall(PetscSFDestroy(&jp->sf));
464 PetscFunctionReturn(PETSC_SUCCESS);
465 }
466
467 /*MC
468 MATCOLORINGJP - Parallel Jones-Plassmann coloring {cite}`jp:pcolor`
469
470 Level: beginner
471
472 Options Database Key:
473 . -mat_coloring_jp_local - perform a local coloring before applying the parallel algorithm
474
475 Notes:
476 This method uses a parallel Luby-style coloring with weights to choose an independent set of processor
477 boundary vertices at each stage that may be assigned colors independently.
478
479 Supports both distance one and distance two colorings.
480
481 .seealso: `MatColoring`, `MatColoringType`, `MatColoringCreate()`, `MatColoring`, `MatColoringSetType()`
482 M*/
MatColoringCreate_JP(MatColoring mc)483 PETSC_EXTERN PetscErrorCode MatColoringCreate_JP(MatColoring mc)
484 {
485 MC_JP *jp;
486
487 PetscFunctionBegin;
488 PetscCall(PetscNew(&jp));
489 jp->sf = NULL;
490 jp->dmask = NULL;
491 jp->omask = NULL;
492 jp->cmask = NULL;
493 jp->dwts = NULL;
494 jp->owts = NULL;
495 jp->local = PETSC_TRUE;
496 mc->data = jp;
497 mc->ops->apply = MatColoringApply_JP;
498 mc->ops->view = NULL;
499 mc->ops->destroy = MatColoringDestroy_JP;
500 mc->ops->setfromoptions = MatColoringSetFromOptions_JP;
501 PetscFunctionReturn(PETSC_SUCCESS);
502 }
503