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