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: // Keep the sign bit clear so mask shifts are defined.
278: maskradix = sizeof(PetscInt) * 8 - 1;
279: maskrounds = 1 + maxcolor / maskradix;
280: maskbase = 0;
281: PetscCall(PetscObjectBaseTypeCompare((PetscObject)G, MATSEQAIJ, &isSeq));
282: PetscCall(PetscObjectTypeCompare((PetscObject)G, MATMPIAIJ, &isMPI));
283: PetscCheck(isSeq || isMPI, PetscObjectComm((PetscObject)G), PETSC_ERR_ARG_WRONGSTATE, "MatColoringDegrees requires an MPI/SEQAIJ Matrix");
285: /* get the inner matrix nonzero structure */
286: oG = NULL;
287: oi = NULL;
288: oj = NULL;
289: if (isMPI) {
290: aij = (Mat_MPIAIJ *)G->data;
291: dG = aij->A;
292: oG = aij->B;
293: daij = (Mat_SeqAIJ *)dG->data;
294: oaij = (Mat_SeqAIJ *)oG->data;
295: di = daij->i;
296: dj = daij->j;
297: oi = oaij->i;
298: oj = oaij->j;
299: PetscCall(MatGetSize(oG, &dn, &on));
300: PetscCall(MatGetMultPetscSF(G, &sf));
301: } else {
302: dG = G;
303: PetscCall(MatGetSize(dG, NULL, &dn));
304: daij = (Mat_SeqAIJ *)dG->data;
305: di = daij->i;
306: dj = daij->j;
307: }
308: for (i = 0; i < dn; i++) mincolors[i] = IS_COLORING_MAX;
309: /* set up the distance-zero mask */
310: if (!dmask) {
311: PetscCall(PetscMalloc1(dn, &dmask));
312: PetscCall(PetscMalloc1(dn, &cmask));
313: jp->cmask = cmask;
314: jp->dmask = dmask;
315: if (oG) {
316: PetscCall(PetscMalloc1(on, &omask));
317: jp->omask = omask;
318: }
319: }
320: /* the number of colors may be more than the number of bits in a PetscInt; take multiple rounds */
321: for (k = 0; k < maskrounds; k++) {
322: for (i = 0; i < dn; i++) {
323: cmask[i] = 0;
324: if (colors[i] < maskbase + maskradix && colors[i] >= maskbase) cmask[i] = ((PetscInt)1) << (colors[i] - maskbase);
325: dmask[i] = cmask[i];
326: }
327: if (oG) {
328: PetscCall(PetscLogEventBegin(MATCOLORING_Comm, mc, 0, 0, 0));
329: PetscCall(PetscSFBcastBegin(sf, MPIU_INT, dmask, omask, MPI_REPLACE));
330: PetscCall(PetscSFBcastEnd(sf, MPIU_INT, dmask, omask, MPI_REPLACE));
331: PetscCall(PetscLogEventEnd(MATCOLORING_Comm, mc, 0, 0, 0));
332: }
333: /* fill in the mask out to the distance of the coloring */
334: for (l = 0; l < mc->dist; l++) {
335: /* fill in the on-and-off diagonal mask */
336: for (i = 0; i < dn; i++) {
337: ncols = di[i + 1] - di[i];
338: cols = &dj[di[i]];
339: for (j = 0; j < ncols; j++) cmask[i] = cmask[i] | dmask[cols[j]];
340: if (oG) {
341: ncols = oi[i + 1] - oi[i];
342: cols = &oj[oi[i]];
343: for (j = 0; j < ncols; j++) cmask[i] = cmask[i] | omask[cols[j]];
344: }
345: }
346: for (i = 0; i < dn; i++) dmask[i] = cmask[i];
347: if (l < mc->dist - 1) {
348: if (oG) {
349: PetscCall(PetscLogEventBegin(MATCOLORING_Comm, mc, 0, 0, 0));
350: PetscCall(PetscSFBcastBegin(sf, MPIU_INT, dmask, omask, MPI_REPLACE));
351: PetscCall(PetscSFBcastEnd(sf, MPIU_INT, dmask, omask, MPI_REPLACE));
352: PetscCall(PetscLogEventEnd(MATCOLORING_Comm, mc, 0, 0, 0));
353: }
354: }
355: }
356: /* read through the mask to see if we've discovered an acceptable color for any vertices in this round */
357: for (i = 0; i < dn; i++) {
358: if (mincolors[i] == IS_COLORING_MAX) {
359: curmask = dmask[i];
360: for (j = 0; j < maskradix; j++) {
361: if (curmask % 2 == 0) {
362: PetscCall(ISColoringValueCast(j + maskbase, &mincolors[i]));
363: break;
364: }
365: curmask = curmask >> 1;
366: }
367: }
368: }
369: /* do the next maskradix colors */
370: maskbase += maskradix;
371: }
372: for (i = 0; i < dn; i++) {
373: if (mincolors[i] == IS_COLORING_MAX) mincolors[i] = maxcolor + 1;
374: }
375: PetscFunctionReturn(PETSC_SUCCESS);
376: }
378: static PetscErrorCode MatColoringApply_JP(MatColoring mc, ISColoring *iscoloring)
379: {
380: MC_JP *jp = (MC_JP *)mc->data;
381: PetscInt i, nadded, nadded_total, nadded_total_old, ntotal, n;
382: PetscInt maxcolor_local = 0, maxcolor_global = 0, *lperm;
383: PetscMPIInt rank;
384: PetscReal *weights, *maxweights;
385: ISColoringValue *color, *mincolor;
387: PetscFunctionBegin;
388: PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mc), &rank));
389: PetscCall(PetscLogEventBegin(MATCOLORING_Weights, mc, 0, 0, 0));
390: PetscCall(MatColoringCreateWeights(mc, &weights, &lperm));
391: PetscCall(PetscLogEventEnd(MATCOLORING_Weights, mc, 0, 0, 0));
392: PetscCall(MatGetSize(mc->mat, NULL, &ntotal));
393: PetscCall(MatGetLocalSize(mc->mat, NULL, &n));
394: PetscCall(PetscMalloc1(n, &maxweights));
395: PetscCall(PetscMalloc1(n, &color));
396: PetscCall(PetscMalloc1(n, &mincolor));
397: for (i = 0; i < n; i++) {
398: color[i] = IS_COLORING_MAX;
399: mincolor[i] = 0;
400: }
401: nadded = 0;
402: nadded_total = 0;
403: nadded_total_old = 0;
404: /* compute purely local vertices */
405: if (jp->local) {
406: PetscCall(MCJPInitialLocalColor_Private(mc, lperm, color));
407: for (i = 0; i < n; i++) {
408: if (color[i] < IS_COLORING_MAX) {
409: nadded++;
410: weights[i] = -1;
411: if (color[i] > maxcolor_local) maxcolor_local = color[i];
412: }
413: }
414: PetscCallMPI(MPIU_Allreduce(&nadded, &nadded_total, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)mc)));
415: PetscCallMPI(MPIU_Allreduce(&maxcolor_local, &maxcolor_global, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)mc)));
416: }
417: while (nadded_total < ntotal) {
418: PetscCall(MCJPMinColor_Private(mc, (ISColoringValue)maxcolor_global, color, mincolor));
419: PetscCall(MCJPGreatestWeight_Private(mc, weights, maxweights));
420: for (i = 0; i < n; i++) {
421: /* choose locally maximal vertices; weights less than zero are omitted from the graph */
422: if (weights[i] >= maxweights[i] && weights[i] >= 0.) {
423: /* assign the minimum possible color */
424: if (mc->maxcolors > mincolor[i]) {
425: color[i] = mincolor[i];
426: } else {
427: color[i] = (ISColoringValue)mc->maxcolors;
428: }
429: if (color[i] > maxcolor_local) maxcolor_local = color[i];
430: weights[i] = -1.;
431: nadded++;
432: }
433: }
434: PetscCallMPI(MPIU_Allreduce(&maxcolor_local, &maxcolor_global, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)mc)));
435: PetscCallMPI(MPIU_Allreduce(&nadded, &nadded_total, 1, MPIU_INT, MPI_SUM, PetscObjectComm((PetscObject)mc)));
436: PetscCheck(nadded_total != nadded_total_old, PetscObjectComm((PetscObject)mc), PETSC_ERR_NOT_CONVERGED, "JP didn't make progress");
437: nadded_total_old = nadded_total;
438: }
439: PetscCall(PetscLogEventBegin(MATCOLORING_ISCreate, mc, 0, 0, 0));
440: PetscCall(ISColoringCreate(PetscObjectComm((PetscObject)mc), maxcolor_global + 1, n, color, PETSC_OWN_POINTER, iscoloring));
441: PetscCall(PetscLogEventEnd(MATCOLORING_ISCreate, mc, 0, 0, 0));
442: PetscCall(PetscFree(jp->dwts));
443: PetscCall(PetscFree(jp->dmask));
444: PetscCall(PetscFree(jp->cmask));
445: PetscCall(PetscFree(jp->owts));
446: PetscCall(PetscFree(jp->omask));
447: PetscCall(PetscFree(weights));
448: PetscCall(PetscFree(lperm));
449: PetscCall(PetscFree(maxweights));
450: PetscCall(PetscFree(mincolor));
451: PetscFunctionReturn(PETSC_SUCCESS);
452: }
454: /*MC
455: MATCOLORINGJP - Parallel Jones-Plassmann coloring {cite}`jp:pcolor`
457: Level: beginner
459: Options Database Key:
460: . -mat_coloring_jp_local - perform a local coloring before applying the parallel algorithm
462: Notes:
463: This method uses a parallel Luby-style coloring with weights to choose an independent set of processor
464: boundary vertices at each stage that may be assigned colors independently.
466: Supports both distance one and distance two colorings.
468: .seealso: [](sec_fdmatrix), [](sec_matfactor), `MatColoringType`, `MatColoring`, `MatColoringCreate()`, `MatColoringSetType()`
469: M*/
470: PETSC_EXTERN PetscErrorCode MatColoringCreate_JP(MatColoring mc)
471: {
472: MC_JP *jp;
474: PetscFunctionBegin;
475: PetscCall(PetscNew(&jp));
476: jp->dmask = NULL;
477: jp->omask = NULL;
478: jp->cmask = NULL;
479: jp->dwts = NULL;
480: jp->owts = NULL;
481: jp->local = PETSC_TRUE;
482: mc->data = jp;
483: mc->ops->apply = MatColoringApply_JP;
484: mc->ops->view = NULL;
485: mc->ops->destroy = MatColoringDestroy_JP;
486: mc->ops->setfromoptions = MatColoringSetFromOptions_JP;
487: PetscFunctionReturn(PETSC_SUCCESS);
488: }