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, in;
    406:   PetscInt               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 (PetscInt 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 (PetscInt 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:       PetscCall(PetscMPIIntCast(n, &in));
    555:       PetscCallMPI(MPI_Allgatherv((void *)is, in, MPIU_INT, cols, ncolsonproc, disp, MPIU_INT, PetscObjectComm((PetscObject)mat)));
    556:     } else if (ctype == IS_COLORING_LOCAL) {
    557:       /* Determine local number of columns of this color on this process, including ghost points */
    558:       nctot = n;
    559:       cols  = (PetscInt *)is;
    560:     } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "Not provided for this MatFDColoring type");

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    726: /*@
    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, 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 (PetscInt 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: }