Actual source code: jp.c
1: #include <../src/mat/impls/aij/mpi/mpiaij.h>
2: #include <petscsf.h>
4: typedef struct {
5: PetscSF sf;
6: PetscReal *dwts, *owts;
7: PetscInt *dmask, *omask, *cmask;
8: PetscBool local;
9: } MC_JP;
11: static PetscErrorCode MatColoringDestroy_JP(MatColoring mc)
12: {
13: PetscFunctionBegin;
14: PetscCall(PetscFree(mc->data));
15: PetscFunctionReturn(PETSC_SUCCESS);
16: }
18: static PetscErrorCode MatColoringSetFromOptions_JP(MatColoring mc, PetscOptionItems *PetscOptionsObject)
19: {
20: MC_JP *jp = (MC_JP *)mc->data;
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: }
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;
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");
50: /* get the inner matrix 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: }
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;
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");
154: /* get the inner matrix 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: }
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;
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");
292: /* get the inner matrix 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: }
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;
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: }
467: /*MC
468: MATCOLORINGJP - Parallel Jones-Plassmann coloring {cite}`jp:pcolor`
470: Level: beginner
472: Options Database Key:
473: . -mat_coloring_jp_local - perform a local coloring before applying the parallel algorithm
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.
479: Supports both distance one and distance two colorings.
481: .seealso: `MatColoring`, `MatColoringType`, `MatColoringCreate()`, `MatColoring`, `MatColoringSetType()`
482: M*/
483: PETSC_EXTERN PetscErrorCode MatColoringCreate_JP(MatColoring mc)
484: {
485: MC_JP *jp;
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: }