Actual source code: fdmpiaij.c

petsc-master 2019-11-16
Report Typos and Errors
  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:   VecPinToCPU(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 instensively in particular piece of scalar CPU code; won't benifit from bouncing back and forth to the GPU */
 65:     VecPinToCPU(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:   VecPinToCPU(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:   VecPinToCPU(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:   VecPinToCPU(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:       VecPinToCPU(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:       VecPinToCPU(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:       VecPinToCPU(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 AD 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: }