Actual source code: fdaij.c

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

      6: /*
      7:     This routine is shared by SeqAIJ and SeqBAIJ matrices,
      8:     since it operators only on the nonzero structure of the elements or blocks.
      9: */
     10: PetscErrorCode MatFDColoringCreate_SeqXAIJ(Mat mat, ISColoring iscoloring, MatFDColoring c)
     11: {
     12:   PetscInt  bs, nis = iscoloring->n, m = mat->rmap->n;
     13:   PetscBool isBAIJ, isSELL;

     15:   PetscFunctionBegin;
     16:   /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian */
     17:   PetscCall(MatGetBlockSize(mat, &bs));
     18:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)mat, MATSEQBAIJ, &isBAIJ));
     19:   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATSEQSELL, &isSELL));
     20:   if (isBAIJ) {
     21:     c->brows = m;
     22:     c->bcols = 1;
     23:   } else { /* seqaij matrix */
     24:     /* bcols is chosen s.t. dy-array takes 50% of memory space as mat */
     25:     PetscReal mem;
     26:     PetscInt  nz, brows, bcols;
     27:     if (isSELL) {
     28:       Mat_SeqSELL *spA = (Mat_SeqSELL *)mat->data;
     29:       nz               = spA->nz;
     30:     } else {
     31:       Mat_SeqAIJ *spA = (Mat_SeqAIJ *)mat->data;
     32:       nz              = spA->nz;
     33:     }

     35:     bs    = 1; /* only bs=1 is supported for SeqAIJ matrix */
     36:     mem   = nz * (sizeof(PetscScalar) + sizeof(PetscInt)) + 3 * m * sizeof(PetscInt);
     37:     bcols = (PetscInt)(0.5 * mem / (m * sizeof(PetscScalar)));
     38:     if (!bcols) bcols = 1;
     39:     brows = 1000 / bcols;
     40:     if (bcols > nis) bcols = nis;
     41:     if (brows == 0 || brows > m) brows = m;
     42:     c->brows = brows;
     43:     c->bcols = bcols;
     44:   }

     46:   c->M       = mat->rmap->N / bs; /* set total rows, columns and local rows */
     47:   c->N       = mat->cmap->N / bs;
     48:   c->m       = mat->rmap->N / bs;
     49:   c->rstart  = 0;
     50:   c->ncolors = nis;
     51:   c->ctype   = iscoloring->ctype;
     52:   PetscFunctionReturn(PETSC_SUCCESS);
     53: }

     55: /*
     56:  Reorder Jentry such that blocked brows*bols of entries from dense matrix are inserted into Jacobian for improved cache performance
     57:    Input Parameters:
     58: +  mat - the matrix containing the nonzero structure of the Jacobian
     59: .  color - the coloring context
     60: -  nz - number of local non-zeros in mat
     61: */
     62: PetscErrorCode MatFDColoringSetUpBlocked_AIJ_Private(Mat mat, MatFDColoring c, PetscInt nz)
     63: {
     64:   PetscInt  i, j, nrows, nbcols, brows = c->brows, bcols = c->bcols, mbs = c->m, nis = c->ncolors;
     65:   PetscInt *color_start, *row_start, *nrows_new, nz_new, row_end;

     67:   PetscFunctionBegin;
     68:   if (brows < 1 || brows > mbs) brows = mbs;
     69:   PetscCall(PetscMalloc2(bcols + 1, &color_start, bcols, &row_start));
     70:   PetscCall(PetscCalloc1(nis, &nrows_new));
     71:   PetscCall(PetscMalloc1(bcols * mat->rmap->n, &c->dy));

     73:   nz_new             = 0;
     74:   nbcols             = 0;
     75:   color_start[bcols] = 0;

     77:   if (c->htype[0] == 'd') { /*  c->htype == 'ds', use MatEntry */
     78:     MatEntry *Jentry_new, *Jentry = c->matentry;

     80:     PetscCall(PetscMalloc1(nz, &Jentry_new));
     81:     for (i = 0; i < nis; i += bcols) { /* loop over colors */
     82:       if (i + bcols > nis) {
     83:         color_start[nis - i] = color_start[bcols];
     84:         bcols                = nis - i;
     85:       }

     87:       color_start[0] = color_start[bcols];
     88:       for (j = 0; j < bcols; j++) {
     89:         color_start[j + 1] = c->nrows[i + j] + color_start[j];
     90:         row_start[j]       = 0;
     91:       }

     93:       row_end = brows;
     94:       if (row_end > mbs) row_end = mbs;

     96:       while (row_end <= mbs) {        /* loop over block rows */
     97:         for (j = 0; j < bcols; j++) { /* loop over block columns */
     98:           nrows = c->nrows[i + j];
     99:           nz    = color_start[j];
    100:           while (row_start[j] < nrows) {
    101:             if (Jentry[nz].row >= row_end) {
    102:               color_start[j] = nz;
    103:               break;
    104:             } else {                                                 /* copy Jentry[nz] to Jentry_new[nz_new] */
    105:               Jentry_new[nz_new].row     = Jentry[nz].row + j * mbs; /* index in dy-array */
    106:               Jentry_new[nz_new].col     = Jentry[nz].col;
    107:               Jentry_new[nz_new].valaddr = Jentry[nz].valaddr;
    108:               nz_new++;
    109:               nz++;
    110:               row_start[j]++;
    111:             }
    112:           }
    113:         }
    114:         if (row_end == mbs) break;
    115:         row_end += brows;
    116:         if (row_end > mbs) row_end = mbs;
    117:       }
    118:       nrows_new[nbcols++] = nz_new;
    119:     }
    120:     PetscCall(PetscFree(Jentry));
    121:     c->matentry = Jentry_new;
    122:   } else { /*  c->htype == 'wp', use MatEntry2 */
    123:     MatEntry2 *Jentry2_new, *Jentry2 = c->matentry2;

    125:     PetscCall(PetscMalloc1(nz, &Jentry2_new));
    126:     for (i = 0; i < nis; i += bcols) { /* loop over colors */
    127:       if (i + bcols > nis) {
    128:         color_start[nis - i] = color_start[bcols];
    129:         bcols                = nis - i;
    130:       }

    132:       color_start[0] = color_start[bcols];
    133:       for (j = 0; j < bcols; j++) {
    134:         color_start[j + 1] = c->nrows[i + j] + color_start[j];
    135:         row_start[j]       = 0;
    136:       }

    138:       row_end = brows;
    139:       if (row_end > mbs) row_end = mbs;

    141:       while (row_end <= mbs) {        /* loop over block rows */
    142:         for (j = 0; j < bcols; j++) { /* loop over block columns */
    143:           nrows = c->nrows[i + j];
    144:           nz    = color_start[j];
    145:           while (row_start[j] < nrows) {
    146:             if (Jentry2[nz].row >= row_end) {
    147:               color_start[j] = nz;
    148:               break;
    149:             } else {                                                   /* copy Jentry2[nz] to Jentry2_new[nz_new] */
    150:               Jentry2_new[nz_new].row     = Jentry2[nz].row + j * mbs; /* index in dy-array */
    151:               Jentry2_new[nz_new].valaddr = Jentry2[nz].valaddr;
    152:               nz_new++;
    153:               nz++;
    154:               row_start[j]++;
    155:             }
    156:           }
    157:         }
    158:         if (row_end == mbs) break;
    159:         row_end += brows;
    160:         if (row_end > mbs) row_end = mbs;
    161:       }
    162:       nrows_new[nbcols++] = nz_new;
    163:     }
    164:     PetscCall(PetscFree(Jentry2));
    165:     c->matentry2 = Jentry2_new;
    166:   } /* ---------------------------------------------*/

    168:   PetscCall(PetscFree2(color_start, row_start));

    170:   for (i = nbcols - 1; i > 0; i--) nrows_new[i] -= nrows_new[i - 1];
    171:   PetscCall(PetscFree(c->nrows));
    172:   c->nrows = nrows_new;
    173:   PetscFunctionReturn(PETSC_SUCCESS);
    174: }

    176: PetscErrorCode MatFDColoringSetUp_SeqXAIJ(Mat mat, ISColoring iscoloring, MatFDColoring c)
    177: {
    178:   PetscInt           i, n, nrows, mbs = c->m, j, k, m, ncols, col, nis = iscoloring->n, *rowhit, bs, bs2, *spidx, nz, tmp, nnz;
    179:   const PetscInt    *is, *row, *ci, *cj;
    180:   PetscBool          isBAIJ, isSELL;
    181:   const PetscScalar *A_val;
    182:   PetscScalar      **valaddrhit;
    183:   MatEntry          *Jentry;
    184:   MatEntry2         *Jentry2;

    186:   PetscFunctionBegin;
    187:   PetscCall(ISColoringGetIS(iscoloring, PETSC_OWN_POINTER, PETSC_IGNORE, &c->isa));

    189:   PetscCall(MatGetBlockSize(mat, &bs));
    190:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)mat, MATSEQBAIJ, &isBAIJ));
    191:   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATSEQSELL, &isSELL));
    192:   if (isBAIJ) {
    193:     Mat_SeqBAIJ *spA = (Mat_SeqBAIJ *)mat->data;

    195:     A_val = spA->a;
    196:     nz    = spA->nz;
    197:   } else if (isSELL) {
    198:     Mat_SeqSELL *spA = (Mat_SeqSELL *)mat->data;

    200:     A_val = spA->val;
    201:     nz    = spA->nz;
    202:     bs    = 1; /* only bs=1 is supported for SeqSELL matrix */
    203:   } else {
    204:     Mat_SeqAIJ *spA = (Mat_SeqAIJ *)mat->data;

    206:     A_val = spA->a;
    207:     nz    = spA->nz;
    208:     bs    = 1; /* only bs=1 is supported for SeqAIJ matrix */
    209:   }
    210:   PetscCall(PetscMalloc2(nis, &c->ncolumns, nis, &c->columns));
    211:   PetscCall(PetscMalloc1(nis, &c->nrows)); /* nrows is freed separately from ncolumns and columns */

    213:   if (c->htype[0] == 'd') {
    214:     PetscCall(PetscMalloc1(nz, &Jentry));
    215:     c->matentry = Jentry;
    216:   } else if (c->htype[0] == 'w') {
    217:     PetscCall(PetscMalloc1(nz, &Jentry2));
    218:     c->matentry2 = Jentry2;
    219:   } else SETERRQ(PetscObjectComm((PetscObject)mat), PETSC_ERR_SUP, "htype is not supported");

    221:   if (isBAIJ) {
    222:     PetscCall(MatGetColumnIJ_SeqBAIJ_Color(mat, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &ci, &cj, &spidx, NULL));
    223:   } else if (isSELL) {
    224:     PetscCall(MatGetColumnIJ_SeqSELL_Color(mat, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &ci, &cj, &spidx, NULL));
    225:   } else {
    226:     PetscCall(MatGetColumnIJ_SeqAIJ_Color(mat, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &ci, &cj, &spidx, NULL));
    227:   }

    229:   PetscCall(PetscCalloc1(c->m, &rowhit));
    230:   PetscCall(PetscMalloc1(c->m, &valaddrhit));

    232:   nnz = 0;
    233:   for (i = 0; i < nis; i++) { /* loop over colors */
    234:     PetscCall(ISGetLocalSize(c->isa[i], &n));
    235:     PetscCall(ISGetIndices(c->isa[i], &is));

    237:     c->ncolumns[i] = n;
    238:     c->columns[i]  = (PetscInt *)is;
    239:     /* note: we know that c->isa is going to be around as long at the c->columns values */
    240:     PetscCall(ISRestoreIndices(c->isa[i], &is));

    242:     /* fast, crude version requires O(N*N) work */
    243:     bs2   = bs * bs;
    244:     nrows = 0;
    245:     for (j = 0; j < n; j++) { /* loop over columns */
    246:       col = is[j];
    247:       tmp = ci[col];
    248:       row = cj + tmp;
    249:       m   = ci[col + 1] - tmp;
    250:       nrows += m;
    251:       for (k = 0; k < m; k++) { /* loop over columns marking them in rowhit */
    252:         PetscAssert(!rowhit[*row], PETSC_COMM_SELF, PETSC_ERR_PLIB, "Incorrect coloring, row %" PetscInt_FMT " shared by multiple columns including column %" PetscInt_FMT " of color %" PetscInt_FMT, col, *row, i);
    253:         rowhit[*row]       = col + 1;
    254:         valaddrhit[*row++] = (PetscScalar *)&A_val[bs2 * spidx[tmp + k]];
    255:       }
    256:     }
    257:     c->nrows[i] = nrows; /* total num of rows for this color */

    259:     if (c->htype[0] == 'd') {
    260:       for (j = 0; j < mbs; j++) { /* loop over rows */
    261:         if (rowhit[j]) {
    262:           Jentry[nnz].row     = j;             /* local row index */
    263:           Jentry[nnz].col     = rowhit[j] - 1; /* local column index */
    264:           Jentry[nnz].valaddr = valaddrhit[j]; /* address of mat value for this entry */
    265:           nnz++;
    266:           rowhit[j] = 0.0; /* zero rowhit for reuse */
    267:         }
    268:       }
    269:     } else {                      /* c->htype == 'wp' */
    270:       for (j = 0; j < mbs; j++) { /* loop over rows */
    271:         if (rowhit[j]) {
    272:           Jentry2[nnz].row     = j;             /* local row index */
    273:           Jentry2[nnz].valaddr = valaddrhit[j]; /* address of mat value for this entry */
    274:           nnz++;
    275:           rowhit[j] = 0.0; /* zero rowhit for reuse */
    276:         }
    277:       }
    278:     }
    279:   }
    280:   PetscCheck(nnz == nz, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Incorrect coloring of matrix");
    281:   if (c->bcols > 1) { /* reorder Jentry for faster MatFDColoringApply() */
    282:     PetscCall(MatFDColoringSetUpBlocked_AIJ_Private(mat, c, nz));
    283:   }

    285:   if (isBAIJ) {
    286:     PetscCall(MatRestoreColumnIJ_SeqBAIJ_Color(mat, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &ci, &cj, &spidx, NULL));
    287:     PetscCall(PetscMalloc1(bs * mat->rmap->n, &c->dy));
    288:   } else if (isSELL) {
    289:     PetscCall(MatRestoreColumnIJ_SeqSELL_Color(mat, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &ci, &cj, &spidx, NULL));
    290:   } else {
    291:     PetscCall(MatRestoreColumnIJ_SeqAIJ_Color(mat, 0, PETSC_FALSE, PETSC_FALSE, &ncols, &ci, &cj, &spidx, NULL));
    292:   }
    293:   PetscCall(PetscFree(rowhit));
    294:   PetscCall(PetscFree(valaddrhit));
    295:   PetscCall(ISColoringRestoreIS(iscoloring, PETSC_OWN_POINTER, &c->isa));

    297:   PetscCall(VecCreateGhost(PetscObjectComm((PetscObject)mat), mat->rmap->n, PETSC_DETERMINE, 0, NULL, &c->vscale));
    298:   PetscCall(PetscInfo(c, "ncolors %" PetscInt_FMT ", brows %" PetscInt_FMT " and bcols %" PetscInt_FMT " are used.\n", c->ncolors, c->brows, c->bcols));
    299:   PetscFunctionReturn(PETSC_SUCCESS);
    300: }