Date: 2022-03-31T04:46:48+00:00

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: PetscErrorCode MatFDColoringApply_BAIJ(Mat J,MatFDColoring coloring,Vec x1,void *sctx)
      7: {
      8:   PetscErrorCode    (*f)(void*,Vec,Vec,void*)=(PetscErrorCode (*)(void*,Vec,Vec,void*))coloring->f;
      9:   PetscInt          k,cstart,cend,l,row,col,nz,spidx,i,j;
     10:   PetscScalar       dx=0.0,*w3_array,*dy_i,*dy=coloring->dy;
     11:   PetscScalar       *vscale_array;
     12:   const PetscScalar *xx;
     13:   PetscReal         epsilon=coloring->error_rel,umin=coloring->umin,unorm;
     14:   Vec               w1=coloring->w1,w2=coloring->w2,w3,vscale=coloring->vscale;
     15:   void              *fctx=coloring->fctx;
     16:   PetscInt          ctype=coloring->ctype,nxloc,nrows_k;
     17:   PetscScalar       *valaddr;
     18:   MatEntry          *Jentry=coloring->matentry;
     19:   MatEntry2         *Jentry2=coloring->matentry2;
     20:   const PetscInt    ncolors=coloring->ncolors,*ncolumns=coloring->ncolumns,*nrows=coloring->nrows;
     21:   PetscInt          bs=J->rmap->bs;

     23:   VecBindToCPU(x1,PETSC_TRUE);
     24:   /* (1) Set w1 = F(x1) */
     25:   if (!coloring->fset) {
     26:     PetscLogEventBegin(MAT_FDColoringFunction,coloring,0,0,0);
     27:     (*f)(sctx,x1,w1,fctx);
     28:     PetscLogEventEnd(MAT_FDColoringFunction,coloring,0,0,0);
     29:   } else {
     30:     coloring->fset = PETSC_FALSE;
     31:   }

     33:   /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */
     34:   VecGetLocalSize(x1,&nxloc);
     35:   if (coloring->htype[0] == 'w') {
     36:     /* vscale = dx is a constant scalar */
     37:     VecNorm(x1,NORM_2,&unorm);
     38:     dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon);
     39:   } else {
     40:     VecGetArrayRead(x1,&xx);
     41:     VecGetArray(vscale,&vscale_array);
     42:     for (col=0; col<nxloc; col++) {
     43:       dx = xx[col];
     44:       if (PetscAbsScalar(dx) < umin) {
     45:         if (PetscRealPart(dx) >= 0.0)      dx = umin;
     46:         else if (PetscRealPart(dx) < 0.0) dx = -umin;
     47:       }
     48:       dx               *= epsilon;
     49:       vscale_array[col] = 1.0/dx;
     50:     }
     51:     VecRestoreArrayRead(x1,&xx);
     52:     VecRestoreArray(vscale,&vscale_array);
     53:   }
     54:   if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') {
     55:     VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD);
     56:     VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD);
     57:   }

     59:   /* (3) Loop over each color */
     60:   if (!coloring->w3) {
     61:     VecDuplicate(x1,&coloring->w3);
     62:     /* Vec is used intensively in particular piece of scalar CPU code; won't benefit from bouncing back and forth to the GPU */
     63:     VecBindToCPU(coloring->w3,PETSC_TRUE);
     64:     PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);
     65:   }
     66:   w3 = coloring->w3;

     68:   VecGetOwnershipRange(x1,&cstart,&cend); /* used by ghosted vscale */
     69:   if (vscale) {
     70:     VecGetArray(vscale,&vscale_array);
     71:   }
     72:   nz = 0;
     73:   for (k=0; k<ncolors; k++) {
     74:     coloring->currentcolor = k;

     76:     /*
     77:       (3-1) Loop over each column associated with color
     78:       adding the perturbation to the vector w3 = x1 + dx.
     79:     */
     80:     VecCopy(x1,w3);
     81:     dy_i = dy;
     82:     for (i=0; i<bs; i++) {     /* Loop over a block of columns */
     83:       VecGetArray(w3,&w3_array);
     84:       if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */
     85:       if (coloring->htype[0] == 'w') {
     86:         for (l=0; l<ncolumns[k]; l++) {
     87:           col            = i + bs*coloring->columns[k][l];  /* local column (in global index!) of the matrix we are probing for */
     88:           w3_array[col] += 1.0/dx;
     89:           if (i) w3_array[col-1] -= 1.0/dx; /* resume original w3[col-1] */
     90:         }
     91:       } else { /* htype == 'ds' */
     92:         vscale_array -= cstart; /* shift pointer so global index can be used */
     93:         for (l=0; l<ncolumns[k]; l++) {
     94:           col = i + bs*coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
     95:           w3_array[col] += 1.0/vscale_array[col];
     96:           if (i) w3_array[col-1] -=  1.0/vscale_array[col-1]; /* resume original w3[col-1] */
     97:         }
     98:         vscale_array += cstart;
     99:       }
    100:       if (ctype == IS_COLORING_GLOBAL) w3_array += cstart;
    101:       VecRestoreArray(w3,&w3_array);

    103:       /*
    104:        (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations)
    105:                            w2 = F(x1 + dx) - F(x1)
    106:        */
    107:       PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);
    108:       VecPlaceArray(w2,dy_i); /* place w2 to the array dy_i */
    109:       (*f)(sctx,w3,w2,fctx);
    110:       PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);
    111:       VecAXPY(w2,-1.0,w1);
    112:       VecResetArray(w2);
    113:       dy_i += nxloc; /* points to dy+i*nxloc */
    114:     }

    116:     /*
    117:      (3-3) Loop over rows of vector, putting results into Jacobian matrix
    118:     */
    119:     nrows_k = nrows[k];
    120:     if (coloring->htype[0] == 'w') {
    121:       for (l=0; l<nrows_k; l++) {
    122:         row     = bs*Jentry2[nz].row;   /* local row index */
    123:         valaddr = Jentry2[nz++].valaddr;
    124:         spidx   = 0;
    125:         dy_i    = dy;
    126:         for (i=0; i<bs; i++) {   /* column of the block */
    127:           for (j=0; j<bs; j++) { /* row of the block */
    128:             valaddr[spidx++] = dy_i[row+j]*dx;
    129:           }
    130:           dy_i += nxloc; /* points to dy+i*nxloc */
    131:         }
    132:       }
    133:     } else { /* htype == 'ds' */
    134:       for (l=0; l<nrows_k; l++) {
    135:         row     = bs*Jentry[nz].row;   /* local row index */
    136:         col     = bs*Jentry[nz].col;   /* local column index */
    137:         valaddr = Jentry[nz++].valaddr;
    138:         spidx   = 0;
    139:         dy_i    = dy;
    140:         for (i=0; i<bs; i++) {   /* column of the block */
    141:           for (j=0; j<bs; j++) { /* row of the block */
    142:             valaddr[spidx++] = dy_i[row+j]*vscale_array[col+i];
    143:           }
    144:           dy_i += nxloc; /* points to dy+i*nxloc */
    145:         }
    146:       }
    147:     }
    148:   }
    149:   MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);
    150:   MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);
    151:   if (vscale) {
    152:     VecRestoreArray(vscale,&vscale_array);
    153:   }

    155:   coloring->currentcolor = -1;
    156:   VecBindToCPU(x1,PETSC_FALSE);
    157:   return 0;
    158: }

    160: /* this is declared PETSC_EXTERN because it is used by MatFDColoringUseDM() which is in the DM library */
    161: PetscErrorCode  MatFDColoringApply_AIJ(Mat J,MatFDColoring coloring,Vec x1,void *sctx)
    162: {
    163:   PetscErrorCode    (*f)(void*,Vec,Vec,void*) = (PetscErrorCode (*)(void*,Vec,Vec,void*))coloring->f;
    164:   PetscInt          k,cstart,cend,l,row,col,nz;
    165:   PetscScalar       dx=0.0,*y,*w3_array;
    166:   const PetscScalar *xx;
    167:   PetscScalar       *vscale_array;
    168:   PetscReal         epsilon=coloring->error_rel,umin=coloring->umin,unorm;
    169:   Vec               w1=coloring->w1,w2=coloring->w2,w3,vscale=coloring->vscale;
    170:   void              *fctx=coloring->fctx;
    171:   ISColoringType    ctype=coloring->ctype;
    172:   PetscInt          nxloc,nrows_k;
    173:   MatEntry          *Jentry=coloring->matentry;
    174:   MatEntry2         *Jentry2=coloring->matentry2;
    175:   const PetscInt    ncolors=coloring->ncolors,*ncolumns=coloring->ncolumns,*nrows=coloring->nrows;
    176:   PetscBool         alreadyboundtocpu;

    178:   VecBoundToCPU(x1,&alreadyboundtocpu);
    179:   VecBindToCPU(x1,PETSC_TRUE);
    181:   /* (1) Set w1 = F(x1) */
    182:   if (!coloring->fset) {
    183:     PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);
    184:     (*f)(sctx,x1,w1,fctx);
    185:     PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);
    186:   } else {
    187:     coloring->fset = PETSC_FALSE;
    188:   }

    190:   /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */
    191:   if (coloring->htype[0] == 'w') {
    192:     /* vscale = 1./dx is a constant scalar */
    193:     VecNorm(x1,NORM_2,&unorm);
    194:     dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon);
    195:   } else {
    196:     VecGetLocalSize(x1,&nxloc);
    197:     VecGetArrayRead(x1,&xx);
    198:     VecGetArray(vscale,&vscale_array);
    199:     for (col=0; col<nxloc; col++) {
    200:       dx = xx[col];
    201:       if (PetscAbsScalar(dx) < umin) {
    202:         if (PetscRealPart(dx) >= 0.0)      dx = umin;
    203:         else if (PetscRealPart(dx) < 0.0) dx = -umin;
    204:       }
    205:       dx               *= epsilon;
    206:       vscale_array[col] = 1.0/dx;
    207:     }
    208:     VecRestoreArrayRead(x1,&xx);
    209:     VecRestoreArray(vscale,&vscale_array);
    210:   }
    211:   if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') {
    212:     VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD);
    213:     VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD);
    214:   }

    216:   /* (3) Loop over each color */
    217:   if (!coloring->w3) {
    218:     VecDuplicate(x1,&coloring->w3);
    219:     PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);
    220:   }
    221:   w3 = coloring->w3;

    223:   VecGetOwnershipRange(x1,&cstart,&cend); /* used by ghosted vscale */
    224:   if (vscale) {
    225:     VecGetArray(vscale,&vscale_array);
    226:   }
    227:   nz = 0;

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

    233:     nbcols = 0;
    234:     for (k=0; k<ncolors; k+=bcols) {

    236:       /*
    237:        (3-1) Loop over each column associated with color
    238:        adding the perturbation to the vector w3 = x1 + dx.
    239:        */

    241:       dy_k = dy;
    242:       if (k + bcols > ncolors) bcols = ncolors - k;
    243:       for (i=0; i<bcols; i++) {
    244:         coloring->currentcolor = k+i;

    246:         VecCopy(x1,w3);
    247:         VecGetArray(w3,&w3_array);
    248:         if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */
    249:         if (coloring->htype[0] == 'w') {
    250:           for (l=0; l<ncolumns[k+i]; l++) {
    251:             col = coloring->columns[k+i][l]; /* local column (in global index!) of the matrix we are probing for */
    252:             w3_array[col] += 1.0/dx;
    253:           }
    254:         } else { /* htype == 'ds' */
    255:           vscale_array -= cstart; /* shift pointer so global index can be used */
    256:           for (l=0; l<ncolumns[k+i]; l++) {
    257:             col = coloring->columns[k+i][l]; /* local column (in global index!) of the matrix we are probing for */
    258:             w3_array[col] += 1.0/vscale_array[col];
    259:           }
    260:           vscale_array += cstart;
    261:         }
    262:         if (ctype == IS_COLORING_GLOBAL) w3_array += cstart;
    263:         VecRestoreArray(w3,&w3_array);

    265:         /*
    266:          (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations)
    267:                            w2 = F(x1 + dx) - F(x1)
    268:          */
    269:         PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);
    270:         VecPlaceArray(w2,dy_k); /* place w2 to the array dy_i */
    271:         (*f)(sctx,w3,w2,fctx);
    272:         PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);
    273:         VecAXPY(w2,-1.0,w1);
    274:         VecResetArray(w2);
    275:         dy_k += m; /* points to dy+i*nxloc */
    276:       }

    278:       /*
    279:        (3-3) Loop over block rows of vector, putting results into Jacobian matrix
    280:        */
    281:       nrows_k = nrows[nbcols++];

    283:       if (coloring->htype[0] == 'w') {
    284:         for (l=0; l<nrows_k; l++) {
    285:           row  = Jentry2[nz].row;   /* local row index */
    286:           /* The 'useless' ifdef is due to a bug in NVIDIA nvc 21.11, which triggers a segfault on this line. We write it in
    287:              another way, and it seems work. See https://lists.mcs.anl.gov/pipermail/petsc-users/2021-December/045158.html
    288:            */
    289:          #if defined(PETSC_USE_COMPLEX)
    290:           PetscScalar *tmp = Jentry2[nz].valaddr;
    291:           *tmp = dy[row]*dx;
    292:          #else
    293:           *(Jentry2[nz].valaddr) = dy[row]*dx;
    294:          #endif
    295:           nz++;
    296:         }
    297:       } else { /* htype == 'ds' */
    298:         for (l=0; l<nrows_k; l++) {
    299:           row = Jentry[nz].row;   /* local row index */
    300:          #if defined(PETSC_USE_COMPLEX) /* See https://lists.mcs.anl.gov/pipermail/petsc-users/2021-December/045158.html */
    301:           PetscScalar *tmp = Jentry[nz].valaddr;
    302:           *tmp = dy[row]*vscale_array[Jentry[nz].col];
    303:          #else
    304:           *(Jentry[nz].valaddr) = dy[row]*vscale_array[Jentry[nz].col];
    305:          #endif
    306:           nz++;
    307:         }
    308:       }
    309:     }
    310:   } else { /* bcols == 1 */
    311:     for (k=0; k<ncolors; k++) {
    312:       coloring->currentcolor = k;

    314:       /*
    315:        (3-1) Loop over each column associated with color
    316:        adding the perturbation to the vector w3 = x1 + dx.
    317:        */
    318:       VecCopy(x1,w3);
    319:       VecGetArray(w3,&w3_array);
    320:       if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */
    321:       if (coloring->htype[0] == 'w') {
    322:         for (l=0; l<ncolumns[k]; l++) {
    323:           col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
    324:           w3_array[col] += 1.0/dx;
    325:         }
    326:       } else { /* htype == 'ds' */
    327:         vscale_array -= cstart; /* shift pointer so global index can be used */
    328:         for (l=0; l<ncolumns[k]; l++) {
    329:           col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
    330:           w3_array[col] += 1.0/vscale_array[col];
    331:         }
    332:         vscale_array += cstart;
    333:       }
    334:       if (ctype == IS_COLORING_GLOBAL) w3_array += cstart;
    335:       VecRestoreArray(w3,&w3_array);

    337:       /*
    338:        (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations)
    339:                            w2 = F(x1 + dx) - F(x1)
    340:        */
    341:       PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);
    342:       (*f)(sctx,w3,w2,fctx);
    343:       PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);
    344:       VecAXPY(w2,-1.0,w1);

    346:       /*
    347:        (3-3) Loop over rows of vector, putting results into Jacobian matrix
    348:        */
    349:       nrows_k = nrows[k];
    350:       VecGetArray(w2,&y);
    351:       if (coloring->htype[0] == 'w') {
    352:         for (l=0; l<nrows_k; l++) {
    353:           row  = Jentry2[nz].row;   /* local row index */
    354:          #if defined(PETSC_USE_COMPLEX) /* See https://lists.mcs.anl.gov/pipermail/petsc-users/2021-December/045158.html */
    355:           PetscScalar *tmp = Jentry2[nz].valaddr;
    356:           *tmp = y[row]*dx;
    357:          #else
    358:           *(Jentry2[nz].valaddr) = y[row]*dx;
    359:          #endif
    360:           nz++;
    361:         }
    362:       } else { /* htype == 'ds' */
    363:         for (l=0; l<nrows_k; l++) {
    364:           row  = Jentry[nz].row;   /* local row index */
    365:          #if defined(PETSC_USE_COMPLEX) /* See https://lists.mcs.anl.gov/pipermail/petsc-users/2021-December/045158.html */
    366:           PetscScalar *tmp = Jentry[nz].valaddr;
    367:           *tmp = y[row]*vscale_array[Jentry[nz].col];
    368:          #else
    369:           *(Jentry[nz].valaddr) = y[row]*vscale_array[Jentry[nz].col];
    370:          #endif
    371:           nz++;
    372:         }
    373:       }
    374:       VecRestoreArray(w2,&y);
    375:     }
    376:   }

    378: #if defined(PETSC_HAVE_VIENNACL) || defined(PETSC_HAVE_CUDA)
    379:   if (J->offloadmask != PETSC_OFFLOAD_UNALLOCATED) J->offloadmask = PETSC_OFFLOAD_CPU;
    380: #endif
    381:   MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);
    382:   MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);
    383:   if (vscale) {
    384:     VecRestoreArray(vscale,&vscale_array);
    385:   }
    386:   coloring->currentcolor = -1;
    387:   if (!alreadyboundtocpu) VecBindToCPU(x1,PETSC_FALSE);
    388:   return 0;
    389: }

    391: PetscErrorCode MatFDColoringSetUp_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
    392: {
    393:   PetscMPIInt            size,*ncolsonproc,*disp,nn;
    394:   PetscInt               i,n,nrows,nrows_i,j,k,m,ncols,col,*rowhit,cstart,cend,colb;
    395:   const PetscInt         *is,*A_ci,*A_cj,*B_ci,*B_cj,*row=NULL,*ltog=NULL;
    396:   PetscInt               nis=iscoloring->n,nctot,*cols,tmp = 0;
    397:   ISLocalToGlobalMapping map=mat->cmap->mapping;
    398:   PetscInt               ctype=c->ctype,*spidxA,*spidxB,nz,bs,bs2,spidx;
    399:   Mat                    A,B;
    400:   PetscScalar            *A_val,*B_val,**valaddrhit;
    401:   MatEntry               *Jentry;
    402:   MatEntry2              *Jentry2;
    403:   PetscBool              isBAIJ,isSELL;
    404:   PetscInt               bcols=c->bcols;
    405: #if defined(PETSC_USE_CTABLE)
    406:   PetscTable             colmap=NULL;
    407: #else
    408:   PetscInt               *colmap=NULL;     /* local col number of off-diag col */
    409: #endif

    411:   if (ctype == IS_COLORING_LOCAL) {
    413:     ISLocalToGlobalMappingGetIndices(map,&ltog);
    414:   }

    416:   MatGetBlockSize(mat,&bs);
    417:   PetscObjectBaseTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);
    418:   PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL);
    419:   if (isBAIJ) {
    420:     Mat_MPIBAIJ *baij=(Mat_MPIBAIJ*)mat->data;
    421:     Mat_SeqBAIJ *spA,*spB;
    422:     A = baij->A;  spA = (Mat_SeqBAIJ*)A->data; A_val = spA->a;
    423:     B = baij->B;  spB = (Mat_SeqBAIJ*)B->data; B_val = spB->a;
    424:     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
    425:     if (!baij->colmap) {
    426:       MatCreateColmap_MPIBAIJ_Private(mat);
    427:     }
    428:     colmap = baij->colmap;
    429:     MatGetColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
    430:     MatGetColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);

    432:     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') {  /* create vscale for storing dx */
    433:       PetscInt    *garray;
    434:       PetscMalloc1(B->cmap->n,&garray);
    435:       for (i=0; i<baij->B->cmap->n/bs; i++) {
    436:         for (j=0; j<bs; j++) {
    437:           garray[i*bs+j] = bs*baij->garray[i]+j;
    438:         }
    439:       }
    440:       VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,garray,&c->vscale);
    441:       VecBindToCPU(c->vscale,PETSC_TRUE);
    442:       PetscFree(garray);
    443:     }
    444:   } else if (isSELL) {
    445:     Mat_MPISELL *sell=(Mat_MPISELL*)mat->data;
    446:     Mat_SeqSELL *spA,*spB;
    447:     A = sell->A;  spA = (Mat_SeqSELL*)A->data; A_val = spA->val;
    448:     B = sell->B;  spB = (Mat_SeqSELL*)B->data; B_val = spB->val;
    449:     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
    450:     if (!sell->colmap) {
    451:       /* Allow access to data structures of local part of matrix
    452:        - creates aij->colmap which maps global column number to local number in part B */
    453:       MatCreateColmap_MPISELL_Private(mat);
    454:     }
    455:     colmap = sell->colmap;
    456:     MatGetColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
    457:     MatGetColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);

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

    461:     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
    462:       VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,sell->garray,&c->vscale);
    463:       VecBindToCPU(c->vscale,PETSC_TRUE);
    464:     }
    465:   } else {
    466:     Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
    467:     Mat_SeqAIJ *spA,*spB;
    468:     A = aij->A;  spA = (Mat_SeqAIJ*)A->data; A_val = spA->a;
    469:     B = aij->B;  spB = (Mat_SeqAIJ*)B->data; B_val = spB->a;
    470:     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
    471:     if (!aij->colmap) {
    472:       /* Allow access to data structures of local part of matrix
    473:        - creates aij->colmap which maps global column number to local number in part B */
    474:       MatCreateColmap_MPIAIJ_Private(mat);
    475:     }
    476:     colmap = aij->colmap;
    477:     MatGetColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
    478:     MatGetColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);

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

    482:     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
    483:       VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,aij->garray,&c->vscale);
    484:       VecBindToCPU(c->vscale,PETSC_TRUE);
    485:     }
    486:   }

    488:   m      = mat->rmap->n/bs;
    489:   cstart = mat->cmap->rstart/bs;
    490:   cend   = mat->cmap->rend/bs;

    492:   PetscMalloc2(nis,&c->ncolumns,nis,&c->columns);
    493:   PetscMalloc1(nis,&c->nrows);
    494:   PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));

    496:   if (c->htype[0] == 'd') {
    497:     PetscMalloc1(nz,&Jentry);
    498:     PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));
    499:     c->matentry = Jentry;
    500:   } else if (c->htype[0] == 'w') {
    501:     PetscMalloc1(nz,&Jentry2);
    502:     PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry2));
    503:     c->matentry2 = Jentry2;
    504:   } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"htype is not supported");

    506:   PetscMalloc2(m+1,&rowhit,m+1,&valaddrhit);
    507:   nz   = 0;
    508:   ISColoringGetIS(iscoloring,PETSC_OWN_POINTER, PETSC_IGNORE,&c->isa);

    510:   if (ctype == IS_COLORING_GLOBAL) {
    511:     MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size);
    512:     PetscMalloc2(size,&ncolsonproc,size,&disp);
    513:   }

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

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

    522:     if (ctype == IS_COLORING_GLOBAL) {
    523:       /* Determine nctot, the total (parallel) number of columns of this color */
    524:       /* ncolsonproc[j]: local ncolumns on proc[j] of this color */
    525:       PetscMPIIntCast(n,&nn);
    526:       MPI_Allgather(&nn,1,MPI_INT,ncolsonproc,1,MPI_INT,PetscObjectComm((PetscObject)mat));
    527:       nctot = 0; for (j=0; j<size; j++) nctot += ncolsonproc[j];
    528:       if (!nctot) {
    529:         PetscInfo(mat,"Coloring of matrix has some unneeded colors with no corresponding rows\n");
    530:       }

    532:       disp[0] = 0;
    533:       for (j=1; j<size; j++) {
    534:         disp[j] = disp[j-1] + ncolsonproc[j-1];
    535:       }

    537:       /* Get cols, the complete list of columns for this color on each process */
    538:       PetscMalloc1(nctot+1,&cols);
    539:       MPI_Allgatherv((void*)is,n,MPIU_INT,cols,ncolsonproc,disp,MPIU_INT,PetscObjectComm((PetscObject)mat));
    540:     } else if (ctype == IS_COLORING_LOCAL) {
    541:       /* Determine local number of columns of this color on this process, including ghost points */
    542:       nctot = n;
    543:       cols  = (PetscInt*)is;
    544:     } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not provided for this MatFDColoring type");

    546:     /* Mark all rows affect by these columns */
    547:     PetscArrayzero(rowhit,m);
    548:     bs2     = bs*bs;
    549:     nrows_i = 0;
    550:     for (j=0; j<nctot; j++) { /* loop over columns*/
    551:       if (ctype == IS_COLORING_LOCAL) {
    552:         col = ltog[cols[j]];
    553:       } else {
    554:         col = cols[j];
    555:       }
    556:       if (col >= cstart && col < cend) { /* column is in A, diagonal block of mat */
    557:         tmp      = A_ci[col-cstart];
    558:         row      = A_cj + tmp;
    559:         nrows    = A_ci[col-cstart+1] - tmp;
    560:         nrows_i += nrows;

    562:         /* loop over columns of A marking them in rowhit */
    563:         for (k=0; k<nrows; k++) {
    564:           /* set valaddrhit for part A */
    565:           spidx            = bs2*spidxA[tmp + k];
    566:           valaddrhit[*row] = &A_val[spidx];
    567:           rowhit[*row++]   = col - cstart + 1; /* local column index */
    568:         }
    569:       } else { /* column is in B, off-diagonal block of mat */
    570: #if defined(PETSC_USE_CTABLE)
    571:         PetscTableFind(colmap,col+1,&colb);
    572:         colb--;
    573: #else
    574:         colb = colmap[col] - 1; /* local column index */
    575: #endif
    576:         if (colb == -1) {
    577:           nrows = 0;
    578:         } else {
    579:           colb  = colb/bs;
    580:           tmp   = B_ci[colb];
    581:           row   = B_cj + tmp;
    582:           nrows = B_ci[colb+1] - tmp;
    583:         }
    584:         nrows_i += nrows;
    585:         /* loop over columns of B marking them in rowhit */
    586:         for (k=0; k<nrows; k++) {
    587:           /* set valaddrhit for part B */
    588:           spidx            = bs2*spidxB[tmp + k];
    589:           valaddrhit[*row] = &B_val[spidx];
    590:           rowhit[*row++]   = colb + 1 + cend - cstart; /* local column index */
    591:         }
    592:       }
    593:     }
    594:     c->nrows[i] = nrows_i;

    596:     if (c->htype[0] == 'd') {
    597:       for (j=0; j<m; j++) {
    598:         if (rowhit[j]) {
    599:           Jentry[nz].row     = j;              /* local row index */
    600:           Jentry[nz].col     = rowhit[j] - 1;  /* local column index */
    601:           Jentry[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
    602:           nz++;
    603:         }
    604:       }
    605:     } else { /* c->htype == 'wp' */
    606:       for (j=0; j<m; j++) {
    607:         if (rowhit[j]) {
    608:           Jentry2[nz].row     = j;              /* local row index */
    609:           Jentry2[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
    610:           nz++;
    611:         }
    612:       }
    613:     }
    614:     if (ctype == IS_COLORING_GLOBAL) {
    615:       PetscFree(cols);
    616:     }
    617:   }
    618:   if (ctype == IS_COLORING_GLOBAL) {
    619:     PetscFree2(ncolsonproc,disp);
    620:   }

    622:   if (bcols > 1) { /* reorder Jentry for faster MatFDColoringApply() */
    623:     MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz);
    624:   }

    626:   if (isBAIJ) {
    627:     MatRestoreColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
    628:     MatRestoreColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);
    629:     PetscMalloc1(bs*mat->rmap->n,&c->dy);
    630:   } else if (isSELL) {
    631:     MatRestoreColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
    632:     MatRestoreColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);
    633:   } else {
    634:     MatRestoreColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
    635:     MatRestoreColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);
    636:   }

    638:   ISColoringRestoreIS(iscoloring,PETSC_OWN_POINTER,&c->isa);
    639:   PetscFree2(rowhit,valaddrhit);

    641:   if (ctype == IS_COLORING_LOCAL) {
    642:     ISLocalToGlobalMappingRestoreIndices(map,&ltog);
    643:   }
    644:   PetscInfo(c,"ncolors %" PetscInt_FMT ", brows %" PetscInt_FMT " and bcols %" PetscInt_FMT " are used.\n",c->ncolors,c->brows,c->bcols);
    645:   return 0;
    646: }

    648: PetscErrorCode MatFDColoringCreate_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
    649: {
    650:   PetscInt       bs,nis=iscoloring->n,m=mat->rmap->n;
    651:   PetscBool      isBAIJ,isSELL;

    653:   /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian;
    654:    bcols is chosen s.t. dy-array takes 50% of memory space as mat */
    655:   MatGetBlockSize(mat,&bs);
    656:   PetscObjectBaseTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);
    657:   PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL);
    658:   if (isBAIJ || m == 0) {
    659:     c->brows = m;
    660:     c->bcols = 1;
    661:   } else if (isSELL) {
    662:     /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
    663:     Mat_MPISELL *sell=(Mat_MPISELL*)mat->data;
    664:     Mat_SeqSELL *spA,*spB;
    665:     Mat        A,B;
    666:     PetscInt   nz,brows,bcols;
    667:     PetscReal  mem;

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

    671:     A = sell->A;  spA = (Mat_SeqSELL*)A->data;
    672:     B = sell->B;  spB = (Mat_SeqSELL*)B->data;
    673:     nz = spA->nz + spB->nz; /* total local nonzero entries of mat */
    674:     mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
    675:     bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
    676:     brows = 1000/bcols;
    677:     if (bcols > nis) bcols = nis;
    678:     if (brows == 0 || brows > m) brows = m;
    679:     c->brows = brows;
    680:     c->bcols = bcols;
    681:   } else { /* mpiaij matrix */
    682:     /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
    683:     Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
    684:     Mat_SeqAIJ *spA,*spB;
    685:     Mat        A,B;
    686:     PetscInt   nz,brows,bcols;
    687:     PetscReal  mem;

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

    691:     A = aij->A;  spA = (Mat_SeqAIJ*)A->data;
    692:     B = aij->B;  spB = (Mat_SeqAIJ*)B->data;
    693:     nz = spA->nz + spB->nz; /* total local nonzero entries of mat */
    694:     mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
    695:     bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
    696:     brows = 1000/bcols;
    697:     if (bcols > nis) bcols = nis;
    698:     if (brows == 0 || brows > m) brows = m;
    699:     c->brows = brows;
    700:     c->bcols = bcols;
    701:   }

    703:   c->M       = mat->rmap->N/bs;         /* set the global rows and columns and local rows */
    704:   c->N       = mat->cmap->N/bs;
    705:   c->m       = mat->rmap->n/bs;
    706:   c->rstart  = mat->rmap->rstart/bs;
    707:   c->ncolors = nis;
    708:   return 0;
    709: }

    711: /*@C

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

    715:    Collective on J

    717:    Input Parameters:
    718: +    J - the sparse matrix
    719: .    coloring - created with MatFDColoringCreate() and a local coloring
    720: -    y - column major storage of matrix values with one color of values per column, the number of rows of y should match
    721:          the number of local rows of J and the number of columns is the number of colors.

    723:    Level: intermediate

    725:    Notes: the matrix in compressed color format may come from an Automatic Differentiation code

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

    729: .seealso: MatFDColoringCreate(), ISColoring, ISColoringCreate(), ISColoringSetType(), IS_COLORING_LOCAL, MatFDColoringSetBlockSize()

    731: @*/
    732: PetscErrorCode  MatFDColoringSetValues(Mat J,MatFDColoring coloring,const PetscScalar *y)
    733: {
    734:   MatEntry2         *Jentry2;
    735:   PetscInt          row,i,nrows_k,l,ncolors,nz = 0,bcols,nbcols = 0;
    736:   const PetscInt    *nrows;
    737:   PetscBool         eq;

    741:   PetscObjectCompareId((PetscObject)J,coloring->matid,&eq);
    743:   Jentry2 = coloring->matentry2;
    744:   nrows   = coloring->nrows;
    745:   ncolors = coloring->ncolors;
    746:   bcols   = coloring->bcols;

    748:   for (i=0; i<ncolors; i+=bcols) {
    749:     nrows_k = nrows[nbcols++];
    750:     for (l=0; l<nrows_k; l++) {
    751:       row                      = Jentry2[nz].row;   /* local row index */
    752:       *(Jentry2[nz++].valaddr) = y[row];
    753:     }
    754:     y += bcols*coloring->m;
    755:   }
    756:   MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);
    757:   MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);
    758:   return 0;
    759: }