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

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

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

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

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

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

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

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

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

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

    181:   VecBindToCPU(x1,PETSC_TRUE);
    182:   if ((ctype == IS_COLORING_LOCAL) && (J->ops->fdcoloringapply == MatFDColoringApply_AIJ)) SETERRQ(PetscObjectComm((PetscObject)J),PETSC_ERR_SUP,"Must call MatColoringUseDM() with IS_COLORING_LOCAL");
    183:   /* (1) Set w1 = F(x1) */
    184:   if (!coloring->fset) {
    185:     PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);
    186:     (*f)(sctx,x1,w1,fctx);
    187:     PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);
    188:   } else {
    189:     coloring->fset = PETSC_FALSE;
    190:   }

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

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

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

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

    235:     nbcols = 0;
    236:     for (k=0; k<ncolors; k+=bcols) {

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

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

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

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

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

    285:       if (coloring->htype[0] == 'w') {
    286:         for (l=0; l<nrows_k; l++) {
    287:           row                      = Jentry2[nz].row;   /* local row index */
    288:           *(Jentry2[nz++].valaddr) = dy[row]*dx;
    289:         }
    290:       } else { /* htype == 'ds' */
    291:         for (l=0; l<nrows_k; l++) {
    292:           row                   = Jentry[nz].row;   /* local row index */
    293:           *(Jentry[nz].valaddr) = dy[row]*vscale_array[Jentry[nz].col];
    294:           nz++;
    295:         }
    296:       }
    297:     }
    298:   } else { /* bcols == 1 */
    299:     for (k=0; k<ncolors; k++) {
    300:       coloring->currentcolor = k;

    302:       /*
    303:        (3-1) Loop over each column associated with color
    304:        adding the perturbation to the vector w3 = x1 + dx.
    305:        */
    306:       VecCopy(x1,w3);
    307:       VecGetArray(w3,&w3_array);
    308:       if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */
    309:       if (coloring->htype[0] == 'w') {
    310:         for (l=0; l<ncolumns[k]; l++) {
    311:           col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
    312:           w3_array[col] += 1.0/dx;
    313:         }
    314:       } else { /* htype == 'ds' */
    315:         vscale_array -= cstart; /* shift pointer so global index can be used */
    316:         for (l=0; l<ncolumns[k]; l++) {
    317:           col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
    318:           w3_array[col] += 1.0/vscale_array[col];
    319:         }
    320:         vscale_array += cstart;
    321:       }
    322:       if (ctype == IS_COLORING_GLOBAL) w3_array += cstart;
    323:       VecRestoreArray(w3,&w3_array);

    325:       /*
    326:        (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations)
    327:                            w2 = F(x1 + dx) - F(x1)
    328:        */
    329:       PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);
    330:       (*f)(sctx,w3,w2,fctx);
    331:       PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);
    332:       VecAXPY(w2,-1.0,w1);

    334:       /*
    335:        (3-3) Loop over rows of vector, putting results into Jacobian matrix
    336:        */
    337:       nrows_k = nrows[k];
    338:       VecGetArray(w2,&y);
    339:       if (coloring->htype[0] == 'w') {
    340:         for (l=0; l<nrows_k; l++) {
    341:           row                      = Jentry2[nz].row;   /* local row index */
    342:           *(Jentry2[nz++].valaddr) = y[row]*dx;
    343:         }
    344:       } else { /* htype == 'ds' */
    345:         for (l=0; l<nrows_k; l++) {
    346:           row                   = Jentry[nz].row;   /* local row index */
    347:           *(Jentry[nz].valaddr) = y[row]*vscale_array[Jentry[nz].col];
    348:           nz++;
    349:         }
    350:       }
    351:       VecRestoreArray(w2,&y);
    352:     }
    353:   }

    355: #if defined(PETSC_HAVE_VIENNACL) || defined(PETSC_HAVE_CUDA)
    356:   if (J->offloadmask != PETSC_OFFLOAD_UNALLOCATED) J->offloadmask = PETSC_OFFLOAD_CPU;
    357: #endif
    358:   MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);
    359:   MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);
    360:   if (vscale) {
    361:     VecRestoreArray(vscale,&vscale_array);
    362:   }
    363:   coloring->currentcolor = -1;
    364:   VecBindToCPU(x1,PETSC_FALSE);
    365:   return(0);
    366: }

    368: PetscErrorCode MatFDColoringSetUp_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
    369: {
    370:   PetscErrorCode         ierr;
    371:   PetscMPIInt            size,*ncolsonproc,*disp,nn;
    372:   PetscInt               i,n,nrows,nrows_i,j,k,m,ncols,col,*rowhit,cstart,cend,colb;
    373:   const PetscInt         *is,*A_ci,*A_cj,*B_ci,*B_cj,*row=NULL,*ltog=NULL;
    374:   PetscInt               nis=iscoloring->n,nctot,*cols,tmp = 0;
    375:   ISLocalToGlobalMapping map=mat->cmap->mapping;
    376:   PetscInt               ctype=c->ctype,*spidxA,*spidxB,nz,bs,bs2,spidx;
    377:   Mat                    A,B;
    378:   PetscScalar            *A_val,*B_val,**valaddrhit;
    379:   MatEntry               *Jentry;
    380:   MatEntry2              *Jentry2;
    381:   PetscBool              isBAIJ,isSELL;
    382:   PetscInt               bcols=c->bcols;
    383: #if defined(PETSC_USE_CTABLE)
    384:   PetscTable             colmap=NULL;
    385: #else
    386:   PetscInt               *colmap=NULL;     /* local col number of off-diag col */
    387: #endif

    390:   if (ctype == IS_COLORING_LOCAL) {
    391:     if (!map) SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_INCOMP,"When using ghosted differencing matrix must have local to global mapping provided with MatSetLocalToGlobalMapping");
    392:     ISLocalToGlobalMappingGetIndices(map,&ltog);
    393:   }

    395:   MatGetBlockSize(mat,&bs);
    396:   PetscObjectBaseTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);
    397:   PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL);
    398:   if (isBAIJ) {
    399:     Mat_MPIBAIJ *baij=(Mat_MPIBAIJ*)mat->data;
    400:     Mat_SeqBAIJ *spA,*spB;
    401:     A = baij->A;  spA = (Mat_SeqBAIJ*)A->data; A_val = spA->a;
    402:     B = baij->B;  spB = (Mat_SeqBAIJ*)B->data; B_val = spB->a;
    403:     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
    404:     if (!baij->colmap) {
    405:       MatCreateColmap_MPIBAIJ_Private(mat);
    406:     }
    407:     colmap = baij->colmap;
    408:     MatGetColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
    409:     MatGetColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);

    411:     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') {  /* create vscale for storing dx */
    412:       PetscInt    *garray;
    413:       PetscMalloc1(B->cmap->n,&garray);
    414:       for (i=0; i<baij->B->cmap->n/bs; i++) {
    415:         for (j=0; j<bs; j++) {
    416:           garray[i*bs+j] = bs*baij->garray[i]+j;
    417:         }
    418:       }
    419:       VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,garray,&c->vscale);
    420:       VecBindToCPU(c->vscale,PETSC_TRUE);
    421:       PetscFree(garray);
    422:     }
    423:   } else if (isSELL) {
    424:     Mat_MPISELL *sell=(Mat_MPISELL*)mat->data;
    425:     Mat_SeqSELL *spA,*spB;
    426:     A = sell->A;  spA = (Mat_SeqSELL*)A->data; A_val = spA->val;
    427:     B = sell->B;  spB = (Mat_SeqSELL*)B->data; B_val = spB->val;
    428:     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
    429:     if (!sell->colmap) {
    430:       /* Allow access to data structures of local part of matrix
    431:        - creates aij->colmap which maps global column number to local number in part B */
    432:       MatCreateColmap_MPISELL_Private(mat);
    433:     }
    434:     colmap = sell->colmap;
    435:     MatGetColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
    436:     MatGetColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);

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

    440:     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
    441:       VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,sell->garray,&c->vscale);
    442:       VecBindToCPU(c->vscale,PETSC_TRUE);
    443:     }
    444:   } else {
    445:     Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
    446:     Mat_SeqAIJ *spA,*spB;
    447:     A = aij->A;  spA = (Mat_SeqAIJ*)A->data; A_val = spA->a;
    448:     B = aij->B;  spB = (Mat_SeqAIJ*)B->data; B_val = spB->a;
    449:     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
    450:     if (!aij->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_MPIAIJ_Private(mat);
    454:     }
    455:     colmap = aij->colmap;
    456:     MatGetColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
    457:     MatGetColumnIJ_SeqAIJ_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,aij->garray,&c->vscale);
    463:       VecBindToCPU(c->vscale,PETSC_TRUE);
    464:     }
    465:   }

    467:   m      = mat->rmap->n/bs;
    468:   cstart = mat->cmap->rstart/bs;
    469:   cend   = mat->cmap->rend/bs;

    471:   PetscMalloc2(nis,&c->ncolumns,nis,&c->columns);
    472:   PetscMalloc1(nis,&c->nrows);
    473:   PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));

    475:   if (c->htype[0] == 'd') {
    476:     PetscMalloc1(nz,&Jentry);
    477:     PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));
    478:     c->matentry = Jentry;
    479:   } else if (c->htype[0] == 'w') {
    480:     PetscMalloc1(nz,&Jentry2);
    481:     PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry2));
    482:     c->matentry2 = Jentry2;
    483:   } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"htype is not supported");

    485:   PetscMalloc2(m+1,&rowhit,m+1,&valaddrhit);
    486:   nz   = 0;
    487:   ISColoringGetIS(iscoloring,PETSC_OWN_POINTER, PETSC_IGNORE,&c->isa);

    489:   if (ctype == IS_COLORING_GLOBAL) {
    490:     MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size);
    491:     PetscMalloc2(size,&ncolsonproc,size,&disp);
    492:   }

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

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

    501:     if (ctype == IS_COLORING_GLOBAL) {
    502:       /* Determine nctot, the total (parallel) number of columns of this color */
    503:       /* ncolsonproc[j]: local ncolumns on proc[j] of this color */
    504:       PetscMPIIntCast(n,&nn);
    505:       MPI_Allgather(&nn,1,MPI_INT,ncolsonproc,1,MPI_INT,PetscObjectComm((PetscObject)mat));
    506:       nctot = 0; for (j=0; j<size; j++) nctot += ncolsonproc[j];
    507:       if (!nctot) {
    508:         PetscInfo(mat,"Coloring of matrix has some unneeded colors with no corresponding rows\n");
    509:       }

    511:       disp[0] = 0;
    512:       for (j=1; j<size; j++) {
    513:         disp[j] = disp[j-1] + ncolsonproc[j-1];
    514:       }

    516:       /* Get cols, the complete list of columns for this color on each process */
    517:       PetscMalloc1(nctot+1,&cols);
    518:       MPI_Allgatherv((void*)is,n,MPIU_INT,cols,ncolsonproc,disp,MPIU_INT,PetscObjectComm((PetscObject)mat));
    519:     } else if (ctype == IS_COLORING_LOCAL) {
    520:       /* Determine local number of columns of this color on this process, including ghost points */
    521:       nctot = n;
    522:       cols  = (PetscInt*)is;
    523:     } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not provided for this MatFDColoring type");

    525:     /* Mark all rows affect by these columns */
    526:     PetscArrayzero(rowhit,m);
    527:     bs2     = bs*bs;
    528:     nrows_i = 0;
    529:     for (j=0; j<nctot; j++) { /* loop over columns*/
    530:       if (ctype == IS_COLORING_LOCAL) {
    531:         col = ltog[cols[j]];
    532:       } else {
    533:         col = cols[j];
    534:       }
    535:       if (col >= cstart && col < cend) { /* column is in A, diagonal block of mat */
    536:         tmp      = A_ci[col-cstart];
    537:         row      = A_cj + tmp;
    538:         nrows    = A_ci[col-cstart+1] - tmp;
    539:         nrows_i += nrows;

    541:         /* loop over columns of A marking them in rowhit */
    542:         for (k=0; k<nrows; k++) {
    543:           /* set valaddrhit for part A */
    544:           spidx            = bs2*spidxA[tmp + k];
    545:           valaddrhit[*row] = &A_val[spidx];
    546:           rowhit[*row++]   = col - cstart + 1; /* local column index */
    547:         }
    548:       } else { /* column is in B, off-diagonal block of mat */
    549: #if defined(PETSC_USE_CTABLE)
    550:         PetscTableFind(colmap,col+1,&colb);
    551:         colb--;
    552: #else
    553:         colb = colmap[col] - 1; /* local column index */
    554: #endif
    555:         if (colb == -1) {
    556:           nrows = 0;
    557:         } else {
    558:           colb  = colb/bs;
    559:           tmp   = B_ci[colb];
    560:           row   = B_cj + tmp;
    561:           nrows = B_ci[colb+1] - tmp;
    562:         }
    563:         nrows_i += nrows;
    564:         /* loop over columns of B marking them in rowhit */
    565:         for (k=0; k<nrows; k++) {
    566:           /* set valaddrhit for part B */
    567:           spidx            = bs2*spidxB[tmp + k];
    568:           valaddrhit[*row] = &B_val[spidx];
    569:           rowhit[*row++]   = colb + 1 + cend - cstart; /* local column index */
    570:         }
    571:       }
    572:     }
    573:     c->nrows[i] = nrows_i;

    575:     if (c->htype[0] == 'd') {
    576:       for (j=0; j<m; j++) {
    577:         if (rowhit[j]) {
    578:           Jentry[nz].row     = j;              /* local row index */
    579:           Jentry[nz].col     = rowhit[j] - 1;  /* local column index */
    580:           Jentry[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
    581:           nz++;
    582:         }
    583:       }
    584:     } else { /* c->htype == 'wp' */
    585:       for (j=0; j<m; j++) {
    586:         if (rowhit[j]) {
    587:           Jentry2[nz].row     = j;              /* local row index */
    588:           Jentry2[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
    589:           nz++;
    590:         }
    591:       }
    592:     }
    593:     if (ctype == IS_COLORING_GLOBAL) {
    594:       PetscFree(cols);
    595:     }
    596:   }
    597:   if (ctype == IS_COLORING_GLOBAL) {
    598:     PetscFree2(ncolsonproc,disp);
    599:   }

    601:   if (bcols > 1) { /* reorder Jentry for faster MatFDColoringApply() */
    602:     MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz);
    603:   }

    605:   if (isBAIJ) {
    606:     MatRestoreColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
    607:     MatRestoreColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);
    608:     PetscMalloc1(bs*mat->rmap->n,&c->dy);
    609:   } else if (isSELL) {
    610:     MatRestoreColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
    611:     MatRestoreColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);
    612:   } else {
    613:     MatRestoreColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
    614:     MatRestoreColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);
    615:   }

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

    620:   if (ctype == IS_COLORING_LOCAL) {
    621:     ISLocalToGlobalMappingRestoreIndices(map,&ltog);
    622:   }
    623:   PetscInfo3(c,"ncolors %D, brows %D and bcols %D are used.\n",c->ncolors,c->brows,c->bcols);
    624:   return(0);
    625: }

    627: PetscErrorCode MatFDColoringCreate_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
    628: {
    630:   PetscInt       bs,nis=iscoloring->n,m=mat->rmap->n;
    631:   PetscBool      isBAIJ,isSELL;

    634:   /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian;
    635:    bcols is chosen s.t. dy-array takes 50% of memory space as mat */
    636:   MatGetBlockSize(mat,&bs);
    637:   PetscObjectBaseTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);
    638:   PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL);
    639:   if (isBAIJ || m == 0) {
    640:     c->brows = m;
    641:     c->bcols = 1;
    642:   } else if (isSELL) {
    643:     /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
    644:     Mat_MPISELL *sell=(Mat_MPISELL*)mat->data;
    645:     Mat_SeqSELL *spA,*spB;
    646:     Mat        A,B;
    647:     PetscInt   nz,brows,bcols;
    648:     PetscReal  mem;

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

    652:     A = sell->A;  spA = (Mat_SeqSELL*)A->data;
    653:     B = sell->B;  spB = (Mat_SeqSELL*)B->data;
    654:     nz = spA->nz + spB->nz; /* total local nonzero entries of mat */
    655:     mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
    656:     bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
    657:     brows = 1000/bcols;
    658:     if (bcols > nis) bcols = nis;
    659:     if (brows == 0 || brows > m) brows = m;
    660:     c->brows = brows;
    661:     c->bcols = bcols;
    662:   } else { /* mpiaij matrix */
    663:     /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
    664:     Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
    665:     Mat_SeqAIJ *spA,*spB;
    666:     Mat        A,B;
    667:     PetscInt   nz,brows,bcols;
    668:     PetscReal  mem;

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

    672:     A = aij->A;  spA = (Mat_SeqAIJ*)A->data;
    673:     B = aij->B;  spB = (Mat_SeqAIJ*)B->data;
    674:     nz = spA->nz + spB->nz; /* total local nonzero entries of mat */
    675:     mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
    676:     bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
    677:     brows = 1000/bcols;
    678:     if (bcols > nis) bcols = nis;
    679:     if (brows == 0 || brows > m) brows = m;
    680:     c->brows = brows;
    681:     c->bcols = bcols;
    682:   }

    684:   c->M       = mat->rmap->N/bs;         /* set the global rows and columns and local rows */
    685:   c->N       = mat->cmap->N/bs;
    686:   c->m       = mat->rmap->n/bs;
    687:   c->rstart  = mat->rmap->rstart/bs;
    688:   c->ncolors = nis;
    689:   return(0);
    690: }

    692: /*@C

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

    696:    Collective on J

    698:    Input Parameters:
    699: +    J - the sparse matrix
    700: .    coloring - created with MatFDColoringCreate() and a local coloring
    701: -    y - column major storage of matrix values with one color of values per column, the number of rows of y should match
    702:          the number of local rows of J and the number of columns is the number of colors.

    704:    Level: intermediate

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

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

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

    712: @*/
    713: PetscErrorCode  MatFDColoringSetValues(Mat J,MatFDColoring coloring,const PetscScalar *y)
    714: {
    715:   PetscErrorCode    ierr;
    716:   MatEntry2         *Jentry2;
    717:   PetscInt          row,i,nrows_k,l,ncolors,nz = 0,bcols,nbcols = 0;
    718:   const PetscInt    *nrows;
    719:   PetscBool         eq;

    724:   PetscObjectCompareId((PetscObject)J,coloring->matid,&eq);
    725:   if (!eq) SETERRQ(PetscObjectComm((PetscObject)J),PETSC_ERR_ARG_WRONG,"Matrix used with MatFDColoringSetValues() must be that used with MatFDColoringCreate()");
    726:   Jentry2 = coloring->matentry2;
    727:   nrows   = coloring->nrows;
    728:   ncolors = coloring->ncolors;
    729:   bcols   = coloring->bcols;

    731:   for (i=0; i<ncolors; i+=bcols) {
    732:     nrows_k = nrows[nbcols++];
    733:     for (l=0; l<nrows_k; l++) {
    734:       row                      = Jentry2[nz].row;   /* local row index */
    735:       *(Jentry2[nz++].valaddr) = y[row];
    736:     }
    737:     y += bcols*coloring->m;
    738:   }
    739:   MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);
    740:   MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);
    741:   return(0);
    742: }