Actual source code: fdmpiaij.c

  1: #include <../src/mat/impls/sell/mpi/mpisell.h>
  2: #include <../src/mat/impls/aij/mpi/mpiaij.h>
  3: #include <../src/mat/impls/baij/mpi/mpibaij.h>
  4: #include <petsc/private/isimpl.h>

  6: static PetscErrorCode MatFDColoringMarkHost_AIJ(Mat J)
  7: {
  8:   PetscBool    isseqAIJ, ismpiAIJ, issell;
  9:   PetscScalar *v;

 11:   PetscFunctionBegin;
 12:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)J, MATMPIAIJ, &ismpiAIJ));
 13:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)J, MATSEQAIJ, &isseqAIJ));
 14:   PetscCall(PetscObjectTypeCompareAny((PetscObject)J, &issell, MATSEQSELLCUDA, MATMPISELLCUDA, ""));
 15:   PetscCheck(!issell, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not coded for %s. Send an email to petsc-dev@mcs.anl.gov to request this feature", ((PetscObject)J)->type_name);
 16:   if (isseqAIJ) {
 17:     PetscCall(MatSeqAIJGetArrayWrite(J, &v));
 18:     PetscCall(MatSeqAIJRestoreArrayWrite(J, &v));
 19:   } else if (ismpiAIJ) {
 20:     Mat dJ, oJ;

 22:     PetscCall(MatMPIAIJGetSeqAIJ(J, &dJ, &oJ, NULL));
 23:     PetscCall(MatSeqAIJGetArrayWrite(dJ, &v));
 24:     PetscCall(MatSeqAIJRestoreArrayWrite(dJ, &v));
 25:     PetscCall(MatSeqAIJGetArrayWrite(oJ, &v));
 26:     PetscCall(MatSeqAIJRestoreArrayWrite(oJ, &v));
 27:   }
 28:   PetscFunctionReturn(PETSC_SUCCESS);
 29: }

 31: PetscErrorCode MatFDColoringApply_BAIJ(Mat J, MatFDColoring coloring, Vec x1, void *sctx)
 32: {
 33:   PetscErrorCode (*f)(void *, Vec, Vec, void *) = (PetscErrorCode(*)(void *, Vec, Vec, void *))coloring->f;
 34:   PetscInt           k, cstart, cend, l, row, col, nz, spidx, i, j;
 35:   PetscScalar        dx = 0.0, *w3_array, *dy_i, *dy = coloring->dy;
 36:   PetscScalar       *vscale_array;
 37:   const PetscScalar *xx;
 38:   PetscReal          epsilon = coloring->error_rel, umin = coloring->umin, unorm;
 39:   Vec                w1 = coloring->w1, w2 = coloring->w2, w3, vscale = coloring->vscale;
 40:   void              *fctx  = coloring->fctx;
 41:   PetscInt           ctype = coloring->ctype, nxloc, nrows_k;
 42:   PetscScalar       *valaddr;
 43:   MatEntry          *Jentry  = coloring->matentry;
 44:   MatEntry2         *Jentry2 = coloring->matentry2;
 45:   const PetscInt     ncolors = coloring->ncolors, *ncolumns = coloring->ncolumns, *nrows = coloring->nrows;
 46:   PetscInt           bs = J->rmap->bs;

 48:   PetscFunctionBegin;
 49:   PetscCall(VecBindToCPU(x1, PETSC_TRUE));
 50:   /* (1) Set w1 = F(x1) */
 51:   if (!coloring->fset) {
 52:     PetscCall(PetscLogEventBegin(MAT_FDColoringFunction, coloring, 0, 0, 0));
 53:     PetscCall((*f)(sctx, x1, w1, fctx));
 54:     PetscCall(PetscLogEventEnd(MAT_FDColoringFunction, coloring, 0, 0, 0));
 55:   } else {
 56:     coloring->fset = PETSC_FALSE;
 57:   }

 59:   /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */
 60:   PetscCall(VecGetLocalSize(x1, &nxloc));
 61:   if (coloring->htype[0] == 'w') {
 62:     /* vscale = dx is a constant scalar */
 63:     PetscCall(VecNorm(x1, NORM_2, &unorm));
 64:     dx = 1.0 / (PetscSqrtReal(1.0 + unorm) * epsilon);
 65:   } else {
 66:     PetscCall(VecGetArrayRead(x1, &xx));
 67:     PetscCall(VecGetArray(vscale, &vscale_array));
 68:     for (col = 0; col < nxloc; col++) {
 69:       dx = xx[col];
 70:       if (PetscAbsScalar(dx) < umin) {
 71:         if (PetscRealPart(dx) >= 0.0) dx = umin;
 72:         else if (PetscRealPart(dx) < 0.0) dx = -umin;
 73:       }
 74:       dx *= epsilon;
 75:       vscale_array[col] = 1.0 / dx;
 76:     }
 77:     PetscCall(VecRestoreArrayRead(x1, &xx));
 78:     PetscCall(VecRestoreArray(vscale, &vscale_array));
 79:   }
 80:   if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') {
 81:     PetscCall(VecGhostUpdateBegin(vscale, INSERT_VALUES, SCATTER_FORWARD));
 82:     PetscCall(VecGhostUpdateEnd(vscale, INSERT_VALUES, SCATTER_FORWARD));
 83:   }

 85:   /* (3) Loop over each color */
 86:   if (!coloring->w3) {
 87:     PetscCall(VecDuplicate(x1, &coloring->w3));
 88:     /* Vec is used intensively in particular piece of scalar CPU code; won't benefit from bouncing back and forth to the GPU */
 89:     PetscCall(VecBindToCPU(coloring->w3, PETSC_TRUE));
 90:   }
 91:   w3 = coloring->w3;

 93:   PetscCall(VecGetOwnershipRange(x1, &cstart, &cend)); /* used by ghosted vscale */
 94:   if (vscale) PetscCall(VecGetArray(vscale, &vscale_array));
 95:   nz = 0;
 96:   for (k = 0; k < ncolors; k++) {
 97:     coloring->currentcolor = k;

 99:     /*
100:       (3-1) Loop over each column associated with color
101:       adding the perturbation to the vector w3 = x1 + dx.
102:     */
103:     PetscCall(VecCopy(x1, w3));
104:     dy_i = dy;
105:     for (i = 0; i < bs; i++) { /* Loop over a block of columns */
106:       PetscCall(VecGetArray(w3, &w3_array));
107:       if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */
108:       if (coloring->htype[0] == 'w') {
109:         for (l = 0; l < ncolumns[k]; l++) {
110:           col = i + bs * coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
111:           w3_array[col] += 1.0 / dx;
112:           if (i) w3_array[col - 1] -= 1.0 / dx; /* resume original w3[col-1] */
113:         }
114:       } else {                  /* htype == 'ds' */
115:         vscale_array -= cstart; /* shift pointer so global index can be used */
116:         for (l = 0; l < ncolumns[k]; l++) {
117:           col = i + bs * coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
118:           w3_array[col] += 1.0 / vscale_array[col];
119:           if (i) w3_array[col - 1] -= 1.0 / vscale_array[col - 1]; /* resume original w3[col-1] */
120:         }
121:         vscale_array += cstart;
122:       }
123:       if (ctype == IS_COLORING_GLOBAL) w3_array += cstart;
124:       PetscCall(VecRestoreArray(w3, &w3_array));

126:       /*
127:        (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations)
128:                            w2 = F(x1 + dx) - F(x1)
129:        */
130:       PetscCall(PetscLogEventBegin(MAT_FDColoringFunction, 0, 0, 0, 0));
131:       PetscCall(VecPlaceArray(w2, dy_i)); /* place w2 to the array dy_i */
132:       PetscCall((*f)(sctx, w3, w2, fctx));
133:       PetscCall(PetscLogEventEnd(MAT_FDColoringFunction, 0, 0, 0, 0));
134:       PetscCall(VecAXPY(w2, -1.0, w1));
135:       PetscCall(VecResetArray(w2));
136:       dy_i += nxloc; /* points to dy+i*nxloc */
137:     }

139:     /*
140:      (3-3) Loop over rows of vector, putting results into Jacobian matrix
141:     */
142:     nrows_k = nrows[k];
143:     if (coloring->htype[0] == 'w') {
144:       for (l = 0; l < nrows_k; l++) {
145:         row     = bs * Jentry2[nz].row; /* local row index */
146:         valaddr = Jentry2[nz++].valaddr;
147:         spidx   = 0;
148:         dy_i    = dy;
149:         for (i = 0; i < bs; i++) {   /* column of the block */
150:           for (j = 0; j < bs; j++) { /* row of the block */
151:             valaddr[spidx++] = dy_i[row + j] * dx;
152:           }
153:           dy_i += nxloc; /* points to dy+i*nxloc */
154:         }
155:       }
156:     } else { /* htype == 'ds' */
157:       for (l = 0; l < nrows_k; l++) {
158:         row     = bs * Jentry[nz].row; /* local row index */
159:         col     = bs * Jentry[nz].col; /* local column index */
160:         valaddr = Jentry[nz++].valaddr;
161:         spidx   = 0;
162:         dy_i    = dy;
163:         for (i = 0; i < bs; i++) {   /* column of the block */
164:           for (j = 0; j < bs; j++) { /* row of the block */
165:             valaddr[spidx++] = dy_i[row + j] * vscale_array[col + i];
166:           }
167:           dy_i += nxloc; /* points to dy+i*nxloc */
168:         }
169:       }
170:     }
171:   }
172:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
173:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
174:   if (vscale) PetscCall(VecRestoreArray(vscale, &vscale_array));

176:   coloring->currentcolor = -1;
177:   PetscCall(VecBindToCPU(x1, PETSC_FALSE));
178:   PetscFunctionReturn(PETSC_SUCCESS);
179: }

181: /* this is declared PETSC_EXTERN because it is used by MatFDColoringUseDM() which is in the DM library */
182: PetscErrorCode MatFDColoringApply_AIJ(Mat J, MatFDColoring coloring, Vec x1, void *sctx)
183: {
184:   PetscErrorCode (*f)(void *, Vec, Vec, void *) = (PetscErrorCode(*)(void *, Vec, Vec, void *))coloring->f;
185:   PetscInt           k, cstart, cend, l, row, col, nz;
186:   PetscScalar        dx = 0.0, *y, *w3_array;
187:   const PetscScalar *xx;
188:   PetscScalar       *vscale_array;
189:   PetscReal          epsilon = coloring->error_rel, umin = coloring->umin, unorm;
190:   Vec                w1 = coloring->w1, w2 = coloring->w2, w3, vscale = coloring->vscale;
191:   void              *fctx  = coloring->fctx;
192:   ISColoringType     ctype = coloring->ctype;
193:   PetscInt           nxloc, nrows_k;
194:   MatEntry          *Jentry  = coloring->matentry;
195:   MatEntry2         *Jentry2 = coloring->matentry2;
196:   const PetscInt     ncolors = coloring->ncolors, *ncolumns = coloring->ncolumns, *nrows = coloring->nrows;
197:   PetscBool          alreadyboundtocpu;

199:   PetscFunctionBegin;
200:   PetscCall(MatFDColoringMarkHost_AIJ(J));
201:   PetscCall(VecBoundToCPU(x1, &alreadyboundtocpu));
202:   PetscCall(VecBindToCPU(x1, PETSC_TRUE));
203:   PetscCheck(!(ctype == IS_COLORING_LOCAL) || !(J->ops->fdcoloringapply == MatFDColoringApply_AIJ), PetscObjectComm((PetscObject)J), PETSC_ERR_SUP, "Must call MatColoringUseDM() with IS_COLORING_LOCAL");
204:   /* (1) Set w1 = F(x1) */
205:   if (!coloring->fset) {
206:     PetscCall(PetscLogEventBegin(MAT_FDColoringFunction, 0, 0, 0, 0));
207:     PetscCall((*f)(sctx, x1, w1, fctx));
208:     PetscCall(PetscLogEventEnd(MAT_FDColoringFunction, 0, 0, 0, 0));
209:   } else {
210:     coloring->fset = PETSC_FALSE;
211:   }

213:   /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */
214:   if (coloring->htype[0] == 'w') {
215:     /* vscale = 1./dx is a constant scalar */
216:     PetscCall(VecNorm(x1, NORM_2, &unorm));
217:     dx = 1.0 / (PetscSqrtReal(1.0 + unorm) * epsilon);
218:   } else {
219:     PetscCall(VecGetLocalSize(x1, &nxloc));
220:     PetscCall(VecGetArrayRead(x1, &xx));
221:     PetscCall(VecGetArray(vscale, &vscale_array));
222:     for (col = 0; col < nxloc; col++) {
223:       dx = xx[col];
224:       if (PetscAbsScalar(dx) < umin) {
225:         if (PetscRealPart(dx) >= 0.0) dx = umin;
226:         else if (PetscRealPart(dx) < 0.0) dx = -umin;
227:       }
228:       dx *= epsilon;
229:       vscale_array[col] = 1.0 / dx;
230:     }
231:     PetscCall(VecRestoreArrayRead(x1, &xx));
232:     PetscCall(VecRestoreArray(vscale, &vscale_array));
233:   }
234:   if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') {
235:     PetscCall(VecGhostUpdateBegin(vscale, INSERT_VALUES, SCATTER_FORWARD));
236:     PetscCall(VecGhostUpdateEnd(vscale, INSERT_VALUES, SCATTER_FORWARD));
237:   }

239:   /* (3) Loop over each color */
240:   if (!coloring->w3) PetscCall(VecDuplicate(x1, &coloring->w3));
241:   w3 = coloring->w3;

243:   PetscCall(VecGetOwnershipRange(x1, &cstart, &cend)); /* used by ghosted vscale */
244:   if (vscale) PetscCall(VecGetArray(vscale, &vscale_array));
245:   nz = 0;

247:   if (coloring->bcols > 1) { /* use blocked insertion of Jentry */
248:     PetscInt     i, m = J->rmap->n, nbcols, bcols = coloring->bcols;
249:     PetscScalar *dy = coloring->dy, *dy_k;

251:     nbcols = 0;
252:     for (k = 0; k < ncolors; k += bcols) {
253:       /*
254:        (3-1) Loop over each column associated with color
255:        adding the perturbation to the vector w3 = x1 + dx.
256:        */

258:       dy_k = dy;
259:       if (k + bcols > ncolors) bcols = ncolors - k;
260:       for (i = 0; i < bcols; i++) {
261:         coloring->currentcolor = k + i;

263:         PetscCall(VecCopy(x1, w3));
264:         PetscCall(VecGetArray(w3, &w3_array));
265:         if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */
266:         if (coloring->htype[0] == 'w') {
267:           for (l = 0; l < ncolumns[k + i]; l++) {
268:             col = coloring->columns[k + i][l]; /* local column (in global index!) of the matrix we are probing for */
269:             w3_array[col] += 1.0 / dx;
270:           }
271:         } else {                  /* htype == 'ds' */
272:           vscale_array -= cstart; /* shift pointer so global index can be used */
273:           for (l = 0; l < ncolumns[k + i]; l++) {
274:             col = coloring->columns[k + i][l]; /* local column (in global index!) of the matrix we are probing for */
275:             w3_array[col] += 1.0 / vscale_array[col];
276:           }
277:           vscale_array += cstart;
278:         }
279:         if (ctype == IS_COLORING_GLOBAL) w3_array += cstart;
280:         PetscCall(VecRestoreArray(w3, &w3_array));

282:         /*
283:          (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations)
284:                            w2 = F(x1 + dx) - F(x1)
285:          */
286:         PetscCall(PetscLogEventBegin(MAT_FDColoringFunction, 0, 0, 0, 0));
287:         PetscCall(VecPlaceArray(w2, dy_k)); /* place w2 to the array dy_i */
288:         PetscCall((*f)(sctx, w3, w2, fctx));
289:         PetscCall(PetscLogEventEnd(MAT_FDColoringFunction, 0, 0, 0, 0));
290:         PetscCall(VecAXPY(w2, -1.0, w1));
291:         PetscCall(VecResetArray(w2));
292:         dy_k += m; /* points to dy+i*nxloc */
293:       }

295:       /*
296:        (3-3) Loop over block rows of vector, putting results into Jacobian matrix
297:        */
298:       nrows_k = nrows[nbcols++];

300:       if (coloring->htype[0] == 'w') {
301:         for (l = 0; l < nrows_k; l++) {
302:           row = Jentry2[nz].row; /* local row index */
303:                                  /* The 'useless' ifdef is due to a bug in NVIDIA nvc 21.11, which triggers a segfault on this line. We write it in
304:              another way, and it seems work. See https://lists.mcs.anl.gov/pipermail/petsc-users/2021-December/045158.html
305:            */
306: #if defined(PETSC_USE_COMPLEX)
307:           PetscScalar *tmp = Jentry2[nz].valaddr;
308:           *tmp             = dy[row] * dx;
309: #else
310:           *(Jentry2[nz].valaddr) = dy[row] * dx;
311: #endif
312:           nz++;
313:         }
314:       } else { /* htype == 'ds' */
315:         for (l = 0; l < nrows_k; l++) {
316:           row = Jentry[nz].row; /* local row index */
317: #if defined(PETSC_USE_COMPLEX)  /* See https://lists.mcs.anl.gov/pipermail/petsc-users/2021-December/045158.html */
318:           PetscScalar *tmp = Jentry[nz].valaddr;
319:           *tmp             = dy[row] * vscale_array[Jentry[nz].col];
320: #else
321:           *(Jentry[nz].valaddr)  = dy[row] * vscale_array[Jentry[nz].col];
322: #endif
323:           nz++;
324:         }
325:       }
326:     }
327:   } else { /* bcols == 1 */
328:     for (k = 0; k < ncolors; k++) {
329:       coloring->currentcolor = k;

331:       /*
332:        (3-1) Loop over each column associated with color
333:        adding the perturbation to the vector w3 = x1 + dx.
334:        */
335:       PetscCall(VecCopy(x1, w3));
336:       PetscCall(VecGetArray(w3, &w3_array));
337:       if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */
338:       if (coloring->htype[0] == 'w') {
339:         for (l = 0; l < ncolumns[k]; l++) {
340:           col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
341:           w3_array[col] += 1.0 / dx;
342:         }
343:       } else {                  /* htype == 'ds' */
344:         vscale_array -= cstart; /* shift pointer so global index can be used */
345:         for (l = 0; l < ncolumns[k]; l++) {
346:           col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
347:           w3_array[col] += 1.0 / vscale_array[col];
348:         }
349:         vscale_array += cstart;
350:       }
351:       if (ctype == IS_COLORING_GLOBAL) w3_array += cstart;
352:       PetscCall(VecRestoreArray(w3, &w3_array));

354:       /*
355:        (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations)
356:                            w2 = F(x1 + dx) - F(x1)
357:        */
358:       PetscCall(PetscLogEventBegin(MAT_FDColoringFunction, 0, 0, 0, 0));
359:       PetscCall((*f)(sctx, w3, w2, fctx));
360:       PetscCall(PetscLogEventEnd(MAT_FDColoringFunction, 0, 0, 0, 0));
361:       PetscCall(VecAXPY(w2, -1.0, w1));

363:       /*
364:        (3-3) Loop over rows of vector, putting results into Jacobian matrix
365:        */
366:       nrows_k = nrows[k];
367:       PetscCall(VecGetArray(w2, &y));
368:       if (coloring->htype[0] == 'w') {
369:         for (l = 0; l < nrows_k; l++) {
370:           row = Jentry2[nz].row; /* local row index */
371: #if defined(PETSC_USE_COMPLEX)   /* See https://lists.mcs.anl.gov/pipermail/petsc-users/2021-December/045158.html */
372:           PetscScalar *tmp = Jentry2[nz].valaddr;
373:           *tmp             = y[row] * dx;
374: #else
375:           *(Jentry2[nz].valaddr) = y[row] * dx;
376: #endif
377:           nz++;
378:         }
379:       } else { /* htype == 'ds' */
380:         for (l = 0; l < nrows_k; l++) {
381:           row = Jentry[nz].row; /* local row index */
382: #if defined(PETSC_USE_COMPLEX)  /* See https://lists.mcs.anl.gov/pipermail/petsc-users/2021-December/045158.html */
383:           PetscScalar *tmp = Jentry[nz].valaddr;
384:           *tmp             = y[row] * vscale_array[Jentry[nz].col];
385: #else
386:           *(Jentry[nz].valaddr)  = y[row] * vscale_array[Jentry[nz].col];
387: #endif
388:           nz++;
389:         }
390:       }
391:       PetscCall(VecRestoreArray(w2, &y));
392:     }
393:   }

395:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
396:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
397:   if (vscale) PetscCall(VecRestoreArray(vscale, &vscale_array));
398:   coloring->currentcolor = -1;
399:   if (!alreadyboundtocpu) PetscCall(VecBindToCPU(x1, PETSC_FALSE));
400:   PetscFunctionReturn(PETSC_SUCCESS);
401: }

403: PetscErrorCode MatFDColoringSetUp_MPIXAIJ(Mat mat, ISColoring iscoloring, MatFDColoring c)
404: {
405:   PetscMPIInt            size, *ncolsonproc, *disp, nn;
406:   PetscInt               i, n, nrows, nrows_i, j, k, m, ncols, col, *rowhit, cstart, cend, colb;
407:   const PetscInt        *is, *A_ci, *A_cj, *B_ci, *B_cj, *row = NULL, *ltog = NULL;
408:   PetscInt               nis = iscoloring->n, nctot, *cols, tmp = 0;
409:   ISLocalToGlobalMapping map   = mat->cmap->mapping;
410:   PetscInt               ctype = c->ctype, *spidxA, *spidxB, nz, bs, bs2, spidx;
411:   Mat                    A, B;
412:   PetscScalar           *A_val, *B_val, **valaddrhit;
413:   MatEntry              *Jentry;
414:   MatEntry2             *Jentry2;
415:   PetscBool              isBAIJ, isSELL;
416:   PetscInt               bcols = c->bcols;
417: #if defined(PETSC_USE_CTABLE)
418:   PetscHMapI colmap = NULL;
419: #else
420:   PetscInt *colmap = NULL;      /* local col number of off-diag col */
421: #endif

423:   PetscFunctionBegin;
424:   if (ctype == IS_COLORING_LOCAL) {
425:     PetscCheck(map, PetscObjectComm((PetscObject)mat), PETSC_ERR_ARG_INCOMP, "When using ghosted differencing matrix must have local to global mapping provided with MatSetLocalToGlobalMapping");
426:     PetscCall(ISLocalToGlobalMappingGetIndices(map, &ltog));
427:   }

429:   PetscCall(MatGetBlockSize(mat, &bs));
430:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)mat, MATMPIBAIJ, &isBAIJ));
431:   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPISELL, &isSELL));
432:   if (isBAIJ) {
433:     Mat_MPIBAIJ *baij = (Mat_MPIBAIJ *)mat->data;
434:     Mat_SeqBAIJ *spA, *spB;
435:     A     = baij->A;
436:     spA   = (Mat_SeqBAIJ *)A->data;
437:     A_val = spA->a;
438:     B     = baij->B;
439:     spB   = (Mat_SeqBAIJ *)B->data;
440:     B_val = spB->a;
441:     nz    = spA->nz + spB->nz; /* total nonzero entries of mat */
442:     if (!baij->colmap) PetscCall(MatCreateColmap_MPIBAIJ_Private(mat));
443:     colmap = baij->colmap;
444:     PetscCall(MatGetColumnIJ_SeqBAIJ_Color(A, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &A_ci, &A_cj, &spidxA, NULL));
445:     PetscCall(MatGetColumnIJ_SeqBAIJ_Color(B, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &B_ci, &B_cj, &spidxB, NULL));

447:     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
448:       PetscInt *garray;
449:       PetscCall(PetscMalloc1(B->cmap->n, &garray));
450:       for (i = 0; i < baij->B->cmap->n / bs; i++) {
451:         for (j = 0; j < bs; j++) garray[i * bs + j] = bs * baij->garray[i] + j;
452:       }
453:       PetscCall(VecCreateGhost(PetscObjectComm((PetscObject)mat), mat->cmap->n, PETSC_DETERMINE, B->cmap->n, garray, &c->vscale));
454:       PetscCall(VecBindToCPU(c->vscale, PETSC_TRUE));
455:       PetscCall(PetscFree(garray));
456:     }
457:   } else if (isSELL) {
458:     Mat_MPISELL *sell = (Mat_MPISELL *)mat->data;
459:     Mat_SeqSELL *spA, *spB;
460:     A     = sell->A;
461:     spA   = (Mat_SeqSELL *)A->data;
462:     A_val = spA->val;
463:     B     = sell->B;
464:     spB   = (Mat_SeqSELL *)B->data;
465:     B_val = spB->val;
466:     nz    = spA->nz + spB->nz; /* total nonzero entries of mat */
467:     if (!sell->colmap) {
468:       /* Allow access to data structures of local part of matrix
469:        - creates aij->colmap which maps global column number to local number in part B */
470:       PetscCall(MatCreateColmap_MPISELL_Private(mat));
471:     }
472:     colmap = sell->colmap;
473:     PetscCall(MatGetColumnIJ_SeqSELL_Color(A, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &A_ci, &A_cj, &spidxA, NULL));
474:     PetscCall(MatGetColumnIJ_SeqSELL_Color(B, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &B_ci, &B_cj, &spidxB, NULL));

476:     bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */

478:     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
479:       PetscCall(VecCreateGhost(PetscObjectComm((PetscObject)mat), mat->cmap->n, PETSC_DETERMINE, B->cmap->n, sell->garray, &c->vscale));
480:       PetscCall(VecBindToCPU(c->vscale, PETSC_TRUE));
481:     }
482:   } else {
483:     Mat_MPIAIJ *aij = (Mat_MPIAIJ *)mat->data;
484:     Mat_SeqAIJ *spA, *spB;
485:     A     = aij->A;
486:     spA   = (Mat_SeqAIJ *)A->data;
487:     A_val = spA->a;
488:     B     = aij->B;
489:     spB   = (Mat_SeqAIJ *)B->data;
490:     B_val = spB->a;
491:     nz    = spA->nz + spB->nz; /* total nonzero entries of mat */
492:     if (!aij->colmap) {
493:       /* Allow access to data structures of local part of matrix
494:        - creates aij->colmap which maps global column number to local number in part B */
495:       PetscCall(MatCreateColmap_MPIAIJ_Private(mat));
496:     }
497:     colmap = aij->colmap;
498:     PetscCall(MatGetColumnIJ_SeqAIJ_Color(A, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &A_ci, &A_cj, &spidxA, NULL));
499:     PetscCall(MatGetColumnIJ_SeqAIJ_Color(B, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &B_ci, &B_cj, &spidxB, NULL));

501:     bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */

503:     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
504:       PetscCall(VecCreateGhost(PetscObjectComm((PetscObject)mat), mat->cmap->n, PETSC_DETERMINE, B->cmap->n, aij->garray, &c->vscale));
505:       PetscCall(VecBindToCPU(c->vscale, PETSC_TRUE));
506:     }
507:   }

509:   m      = mat->rmap->n / bs;
510:   cstart = mat->cmap->rstart / bs;
511:   cend   = mat->cmap->rend / bs;

513:   PetscCall(PetscMalloc2(nis, &c->ncolumns, nis, &c->columns));
514:   PetscCall(PetscMalloc1(nis, &c->nrows));

516:   if (c->htype[0] == 'd') {
517:     PetscCall(PetscMalloc1(nz, &Jentry));
518:     c->matentry = Jentry;
519:   } else if (c->htype[0] == 'w') {
520:     PetscCall(PetscMalloc1(nz, &Jentry2));
521:     c->matentry2 = Jentry2;
522:   } else SETERRQ(PetscObjectComm((PetscObject)mat), PETSC_ERR_SUP, "htype is not supported");

524:   PetscCall(PetscMalloc2(m + 1, &rowhit, m + 1, &valaddrhit));
525:   nz = 0;
526:   PetscCall(ISColoringGetIS(iscoloring, PETSC_OWN_POINTER, PETSC_IGNORE, &c->isa));

528:   if (ctype == IS_COLORING_GLOBAL) {
529:     PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size));
530:     PetscCall(PetscMalloc2(size, &ncolsonproc, size, &disp));
531:   }

533:   for (i = 0; i < nis; i++) { /* for each local color */
534:     PetscCall(ISGetLocalSize(c->isa[i], &n));
535:     PetscCall(ISGetIndices(c->isa[i], &is));

537:     c->ncolumns[i] = n; /* local number of columns of this color on this process */
538:     c->columns[i]  = (PetscInt *)is;

540:     if (ctype == IS_COLORING_GLOBAL) {
541:       /* Determine nctot, the total (parallel) number of columns of this color */
542:       /* ncolsonproc[j]: local ncolumns on proc[j] of this color */
543:       PetscCall(PetscMPIIntCast(n, &nn));
544:       PetscCallMPI(MPI_Allgather(&nn, 1, MPI_INT, ncolsonproc, 1, MPI_INT, PetscObjectComm((PetscObject)mat)));
545:       nctot = 0;
546:       for (j = 0; j < size; j++) nctot += ncolsonproc[j];
547:       if (!nctot) PetscCall(PetscInfo(mat, "Coloring of matrix has some unneeded colors with no corresponding rows\n"));

549:       disp[0] = 0;
550:       for (j = 1; j < size; j++) disp[j] = disp[j - 1] + ncolsonproc[j - 1];

552:       /* Get cols, the complete list of columns for this color on each process */
553:       PetscCall(PetscMalloc1(nctot + 1, &cols));
554:       PetscCallMPI(MPI_Allgatherv((void *)is, n, MPIU_INT, cols, ncolsonproc, disp, MPIU_INT, PetscObjectComm((PetscObject)mat)));
555:     } else if (ctype == IS_COLORING_LOCAL) {
556:       /* Determine local number of columns of this color on this process, including ghost points */
557:       nctot = n;
558:       cols  = (PetscInt *)is;
559:     } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "Not provided for this MatFDColoring type");

561:     /* Mark all rows affect by these columns */
562:     PetscCall(PetscArrayzero(rowhit, m));
563:     bs2     = bs * bs;
564:     nrows_i = 0;
565:     for (j = 0; j < nctot; j++) { /* loop over columns*/
566:       if (ctype == IS_COLORING_LOCAL) {
567:         col = ltog[cols[j]];
568:       } else {
569:         col = cols[j];
570:       }
571:       if (col >= cstart && col < cend) { /* column is in A, diagonal block of mat */
572:         tmp   = A_ci[col - cstart];
573:         row   = A_cj + tmp;
574:         nrows = A_ci[col - cstart + 1] - tmp;
575:         nrows_i += nrows;

577:         /* loop over columns of A marking them in rowhit */
578:         for (k = 0; k < nrows; k++) {
579:           /* set valaddrhit for part A */
580:           spidx            = bs2 * spidxA[tmp + k];
581:           valaddrhit[*row] = &A_val[spidx];
582:           rowhit[*row++]   = col - cstart + 1; /* local column index */
583:         }
584:       } else { /* column is in B, off-diagonal block of mat */
585: #if defined(PETSC_USE_CTABLE)
586:         PetscCall(PetscHMapIGetWithDefault(colmap, col + 1, 0, &colb));
587:         colb--;
588: #else
589:         colb = colmap[col] - 1; /* local column index */
590: #endif
591:         if (colb == -1) {
592:           nrows = 0;
593:         } else {
594:           colb  = colb / bs;
595:           tmp   = B_ci[colb];
596:           row   = B_cj + tmp;
597:           nrows = B_ci[colb + 1] - tmp;
598:         }
599:         nrows_i += nrows;
600:         /* loop over columns of B marking them in rowhit */
601:         for (k = 0; k < nrows; k++) {
602:           /* set valaddrhit for part B */
603:           spidx            = bs2 * spidxB[tmp + k];
604:           valaddrhit[*row] = &B_val[spidx];
605:           rowhit[*row++]   = colb + 1 + cend - cstart; /* local column index */
606:         }
607:       }
608:     }
609:     c->nrows[i] = nrows_i;

611:     if (c->htype[0] == 'd') {
612:       for (j = 0; j < m; j++) {
613:         if (rowhit[j]) {
614:           Jentry[nz].row     = j;             /* local row index */
615:           Jentry[nz].col     = rowhit[j] - 1; /* local column index */
616:           Jentry[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */
617:           nz++;
618:         }
619:       }
620:     } else { /* c->htype == 'wp' */
621:       for (j = 0; j < m; j++) {
622:         if (rowhit[j]) {
623:           Jentry2[nz].row     = j;             /* local row index */
624:           Jentry2[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */
625:           nz++;
626:         }
627:       }
628:     }
629:     if (ctype == IS_COLORING_GLOBAL) PetscCall(PetscFree(cols));
630:   }
631:   if (ctype == IS_COLORING_GLOBAL) PetscCall(PetscFree2(ncolsonproc, disp));

633:   if (bcols > 1) { /* reorder Jentry for faster MatFDColoringApply() */
634:     PetscCall(MatFDColoringSetUpBlocked_AIJ_Private(mat, c, nz));
635:   }

637:   if (isBAIJ) {
638:     PetscCall(MatRestoreColumnIJ_SeqBAIJ_Color(A, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &A_ci, &A_cj, &spidxA, NULL));
639:     PetscCall(MatRestoreColumnIJ_SeqBAIJ_Color(B, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &B_ci, &B_cj, &spidxB, NULL));
640:     PetscCall(PetscMalloc1(bs * mat->rmap->n, &c->dy));
641:   } else if (isSELL) {
642:     PetscCall(MatRestoreColumnIJ_SeqSELL_Color(A, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &A_ci, &A_cj, &spidxA, NULL));
643:     PetscCall(MatRestoreColumnIJ_SeqSELL_Color(B, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &B_ci, &B_cj, &spidxB, NULL));
644:   } else {
645:     PetscCall(MatRestoreColumnIJ_SeqAIJ_Color(A, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &A_ci, &A_cj, &spidxA, NULL));
646:     PetscCall(MatRestoreColumnIJ_SeqAIJ_Color(B, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &B_ci, &B_cj, &spidxB, NULL));
647:   }

649:   PetscCall(ISColoringRestoreIS(iscoloring, PETSC_OWN_POINTER, &c->isa));
650:   PetscCall(PetscFree2(rowhit, valaddrhit));

652:   if (ctype == IS_COLORING_LOCAL) PetscCall(ISLocalToGlobalMappingRestoreIndices(map, &ltog));
653:   PetscCall(PetscInfo(c, "ncolors %" PetscInt_FMT ", brows %" PetscInt_FMT " and bcols %" PetscInt_FMT " are used.\n", c->ncolors, c->brows, c->bcols));
654:   PetscFunctionReturn(PETSC_SUCCESS);
655: }

657: PetscErrorCode MatFDColoringCreate_MPIXAIJ(Mat mat, ISColoring iscoloring, MatFDColoring c)
658: {
659:   PetscInt  bs, nis = iscoloring->n, m = mat->rmap->n;
660:   PetscBool isBAIJ, isSELL;

662:   PetscFunctionBegin;
663:   /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian;
664:    bcols is chosen s.t. dy-array takes 50% of memory space as mat */
665:   PetscCall(MatGetBlockSize(mat, &bs));
666:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)mat, MATMPIBAIJ, &isBAIJ));
667:   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPISELL, &isSELL));
668:   if (isBAIJ || m == 0) {
669:     c->brows = m;
670:     c->bcols = 1;
671:   } else if (isSELL) {
672:     /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
673:     Mat_MPISELL *sell = (Mat_MPISELL *)mat->data;
674:     Mat_SeqSELL *spA, *spB;
675:     Mat          A, B;
676:     PetscInt     nz, brows, bcols;
677:     PetscReal    mem;

679:     bs = 1; /* only bs=1 is supported for MPISELL matrix */

681:     A     = sell->A;
682:     spA   = (Mat_SeqSELL *)A->data;
683:     B     = sell->B;
684:     spB   = (Mat_SeqSELL *)B->data;
685:     nz    = spA->nz + spB->nz; /* total local nonzero entries of mat */
686:     mem   = nz * (sizeof(PetscScalar) + sizeof(PetscInt)) + 3 * m * sizeof(PetscInt);
687:     bcols = (PetscInt)(0.5 * mem / (m * sizeof(PetscScalar)));
688:     brows = 1000 / bcols;
689:     if (bcols > nis) bcols = nis;
690:     if (brows == 0 || brows > m) brows = m;
691:     c->brows = brows;
692:     c->bcols = bcols;
693:   } else { /* mpiaij matrix */
694:     /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
695:     Mat_MPIAIJ *aij = (Mat_MPIAIJ *)mat->data;
696:     Mat_SeqAIJ *spA, *spB;
697:     Mat         A, B;
698:     PetscInt    nz, brows, bcols;
699:     PetscReal   mem;

701:     bs = 1; /* only bs=1 is supported for MPIAIJ matrix */

703:     A     = aij->A;
704:     spA   = (Mat_SeqAIJ *)A->data;
705:     B     = aij->B;
706:     spB   = (Mat_SeqAIJ *)B->data;
707:     nz    = spA->nz + spB->nz; /* total local nonzero entries of mat */
708:     mem   = nz * (sizeof(PetscScalar) + sizeof(PetscInt)) + 3 * m * sizeof(PetscInt);
709:     bcols = (PetscInt)(0.5 * mem / (m * sizeof(PetscScalar)));
710:     brows = 1000 / bcols;
711:     if (bcols > nis) bcols = nis;
712:     if (brows == 0 || brows > m) brows = m;
713:     c->brows = brows;
714:     c->bcols = bcols;
715:   }

717:   c->M       = mat->rmap->N / bs; /* set the global rows and columns and local rows */
718:   c->N       = mat->cmap->N / bs;
719:   c->m       = mat->rmap->n / bs;
720:   c->rstart  = mat->rmap->rstart / bs;
721:   c->ncolors = nis;
722:   PetscFunctionReturn(PETSC_SUCCESS);
723: }

725: /*@C

727:   MatFDColoringSetValues - takes a matrix in compressed color format and enters the matrix into a PETSc `Mat`

729:   Collective

731:   Input Parameters:
732: + J        - the sparse matrix
733: . coloring - created with `MatFDColoringCreate()` and a local coloring
734: - y        - column major storage of matrix values with one color of values per column, the number of rows of y should match
735:          the number of local rows of `J` and the number of columns is the number of colors.

737:   Level: intermediate

739:   Notes:
740:   The matrix in compressed color format may come from an automatic differentiation code

742:   The code will be slightly faster if `MatFDColoringSetBlockSize`(coloring,`PETSC_DEFAULT`,nc); is called immediately after creating the coloring

744: .seealso: [](ch_matrices), `Mat`, `MatFDColoringCreate()`, `ISColoring`, `ISColoringCreate()`, `ISColoringSetType()`, `IS_COLORING_LOCAL`, `MatFDColoringSetBlockSize()`
745: @*/
746: PetscErrorCode MatFDColoringSetValues(Mat J, MatFDColoring coloring, const PetscScalar *y)
747: {
748:   MatEntry2      *Jentry2;
749:   PetscInt        row, i, nrows_k, l, ncolors, nz = 0, bcols, nbcols = 0;
750:   const PetscInt *nrows;
751:   PetscBool       eq;

753:   PetscFunctionBegin;
756:   PetscCall(PetscObjectCompareId((PetscObject)J, coloring->matid, &eq));
757:   PetscCheck(eq, PetscObjectComm((PetscObject)J), PETSC_ERR_ARG_WRONG, "Matrix used with MatFDColoringSetValues() must be that used with MatFDColoringCreate()");
758:   Jentry2 = coloring->matentry2;
759:   nrows   = coloring->nrows;
760:   ncolors = coloring->ncolors;
761:   bcols   = coloring->bcols;

763:   for (i = 0; i < ncolors; i += bcols) {
764:     nrows_k = nrows[nbcols++];
765:     for (l = 0; l < nrows_k; l++) {
766:       row                      = Jentry2[nz].row; /* local row index */
767:       *(Jentry2[nz++].valaddr) = y[row];
768:     }
769:     y += bcols * coloring->m;
770:   }
771:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
772:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
773:   PetscFunctionReturn(PETSC_SUCCESS);
774: }