Actual source code: fdaij.c

petsc-3.11.4 2019-09-28
Report Typos and Errors
  1:  #include <../src/mat/impls/aij/seq/aij.h>
  2:  #include <../src/mat/impls/baij/seq/baij.h>
  3:  #include <../src/mat/impls/sell/seq/sell.h>
  4:  #include <petsc/private/isimpl.h>

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

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

 36:     bs    = 1; /* only bs=1 is supported for SeqAIJ matrix */

 38:     mem   = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
 39:     bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
 40:     brows = 1000/bcols;
 41:     if (bcols > nis) bcols = nis;
 42:     if (brows == 0 || brows > m) brows = m;
 43:     c->brows = brows;
 44:     c->bcols = bcols;
 45:   }

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

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

 70:   if (brows < 1 || brows > mbs) brows = mbs;
 71:   PetscMalloc2(bcols+1,&color_start,bcols,&row_start);
 72:   PetscCalloc1(nis,&nrows_new);
 73:   PetscMalloc1(bcols*mat->rmap->n,&c->dy);
 74:   PetscLogObjectMemory((PetscObject)c,bcols*mat->rmap->n*sizeof(PetscScalar));

 76:   nz_new = 0;
 77:   nbcols = 0;
 78:   color_start[bcols] = 0;

 80:   if (c->htype[0] == 'd') { /* ----  c->htype == 'ds', use MatEntry --------*/
 81:     MatEntry       *Jentry_new,*Jentry=c->matentry;
 82:     PetscMalloc1(nz,&Jentry_new);
 83:     for (i=0; i<nis; i+=bcols) { /* loop over colors */
 84:       if (i + bcols > nis) {
 85:         color_start[nis - i] = color_start[bcols];
 86:         bcols                = nis - i;
 87:       }

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

 95:       row_end = brows;
 96:       if (row_end > mbs) row_end = mbs;

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

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

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

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

165:   PetscFree2(color_start,row_start);

167:   for (i=nbcols-1; i>0; i--) nrows_new[i] -= nrows_new[i-1];
168:   PetscFree(c->nrows);
169:   c->nrows = nrows_new;
170:   return(0);
171: }

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

185:   ISColoringGetIS(iscoloring,PETSC_IGNORE,&isa);

187:   MatGetBlockSize(mat,&bs);
188:   PetscObjectTypeCompare((PetscObject)mat,MATSEQBAIJ,&isBAIJ);
189:   PetscObjectTypeCompare((PetscObject)mat,MATSEQSELL,&isSELL);
190:   if (isBAIJ) {
191:     Mat_SeqBAIJ *spA = (Mat_SeqBAIJ*)mat->data;
192:     A_val = spA->a;
193:     nz    = spA->nz;
194:   } else if (isSELL) {
195:     Mat_SeqSELL  *spA = (Mat_SeqSELL*)mat->data;
196:     A_val = spA->val;
197:     nz    = spA->nz;
198:     bs    = 1; /* only bs=1 is supported for SeqSELL matrix */
199:   } else {
200:     Mat_SeqAIJ  *spA = (Mat_SeqAIJ*)mat->data;
201:     A_val = spA->a;
202:     nz    = spA->nz;
203:     bs    = 1; /* only bs=1 is supported for SeqAIJ matrix */
204:   }

206:   PetscMalloc1(nis,&c->ncolumns);
207:   PetscMalloc1(nis,&c->columns);
208:   PetscMalloc1(nis,&c->nrows);
209:   PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));

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

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

229:   PetscMalloc2(c->m,&rowhit,c->m,&valaddrhit);
230:   PetscMemzero(rowhit,c->m*sizeof(PetscInt));

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

237:     c->ncolumns[i] = n;
238:     if (n) {
239:       PetscMalloc1(n,&c->columns[i]);
240:       PetscLogObjectMemory((PetscObject)c,n*sizeof(PetscInt));
241:       PetscMemcpy(c->columns[i],is,n*sizeof(PetscInt));
242:     } else {
243:       c->columns[i] = 0;
244:     }

246:     /* fast, crude version requires O(N*N) work */
247:     bs2   = bs*bs;
248:     nrows = 0;
249:     for (j=0; j<n; j++) {  /* loop over columns */
250:       col    = is[j];
251:       row    = cj + ci[col];
252:       m      = ci[col+1] - ci[col];
253:       nrows += m;
254:       for (k=0; k<m; k++) {  /* loop over columns marking them in rowhit */
255:         rowhit[*row]       = col + 1;
256:         valaddrhit[*row++] = &A_val[bs2*spidx[ci[col] + k]];
257:       }
258:     }
259:     c->nrows[i] = nrows; /* total num of rows for this color */

261:     if (c->htype[0] == 'd') {
262:       for (j=0; j<mbs; j++) { /* loop over rows */
263:         if (rowhit[j]) {
264:           Jentry[nz].row     = j;              /* local row index */
265:           Jentry[nz].col     = rowhit[j] - 1;  /* local column index */
266:           Jentry[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
267:           nz++;
268:           rowhit[j] = 0.0;                     /* zero rowhit for reuse */
269:         }
270:       }
271:     }  else { /* c->htype == 'wp' */
272:       for (j=0; j<mbs; j++) { /* loop over rows */
273:         if (rowhit[j]) {
274:           Jentry2[nz].row     = j;              /* local row index */
275:           Jentry2[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
276:           nz++;
277:           rowhit[j] = 0.0;                     /* zero rowhit for reuse */
278:         }
279:       }
280:     }
281:     ISRestoreIndices(isa[i],&is);
282:   }

284:   if (c->bcols > 1) {  /* reorder Jentry for faster MatFDColoringApply() */
285:     MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz);
286:   }

288:   if (isBAIJ) {
289:     MatRestoreColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);
290:     PetscMalloc1(bs*mat->rmap->n,&c->dy);
291:     PetscLogObjectMemory((PetscObject)c,bs*mat->rmap->n*sizeof(PetscScalar));
292:   } else if (isSELL) {
293:     MatRestoreColumnIJ_SeqSELL_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);
294:   } else {
295:     MatRestoreColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);
296:   }
297:   PetscFree2(rowhit,valaddrhit);
298:   ISColoringRestoreIS(iscoloring,&isa);

300:   VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->rmap->n,PETSC_DETERMINE,0,NULL,&c->vscale);
301:   PetscInfo3(c,"ncolors %D, brows %D and bcols %D are used.\n",c->ncolors,c->brows,c->bcols);
302:   return(0);
303: }