xref: /petsc/src/mat/impls/aij/seq/fdaij.c (revision 5b6bfdb9644f185dbf5e5a09b808ec241507e1e7)
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>
5 
6 /*
7     This routine is shared by SeqAIJ and SeqBAIJ matrices,
8     since it operators only on the nonzero structure of the elements or blocks.
9 */
10 PetscErrorCode MatFDColoringCreate_SeqXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
11 {
12   PetscErrorCode ierr;
13   PetscInt       bs,nis=iscoloring->n,m=mat->rmap->n;
14   PetscBool      isBAIJ,isSELL;
15 
16   PetscFunctionBegin;
17   /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian */
18   ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
19   ierr = PetscObjectTypeCompare((PetscObject)mat,MATSEQBAIJ,&isBAIJ);CHKERRQ(ierr);
20   ierr = PetscObjectTypeCompare((PetscObject)mat,MATSEQSELL,&isSELL);CHKERRQ(ierr);
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     }
35 
36     bs    = 1; /* only bs=1 is supported for SeqAIJ matrix */
37 
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   }
46 
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   PetscFunctionReturn(0);
54 }
55 
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 {
65   PetscErrorCode ierr;
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;
68 
69   PetscFunctionBegin;
70   if (brows < 1 || brows > mbs) brows = mbs;
71   ierr = PetscMalloc2(bcols+1,&color_start,bcols,&row_start);CHKERRQ(ierr);
72   ierr = PetscCalloc1(nis,&nrows_new);CHKERRQ(ierr);
73   ierr = PetscMalloc1(bcols*mat->rmap->n,&c->dy);CHKERRQ(ierr);
74   ierr = PetscLogObjectMemory((PetscObject)c,bcols*mat->rmap->n*sizeof(PetscScalar));CHKERRQ(ierr);
75 
76   nz_new = 0;
77   nbcols = 0;
78   color_start[bcols] = 0;
79 
80   if (c->htype[0] == 'd') { /* ----  c->htype == 'ds', use MatEntry --------*/
81     MatEntry       *Jentry_new,*Jentry=c->matentry;
82     ierr = PetscMalloc1(nz,&Jentry_new);CHKERRQ(ierr);
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       }
88 
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       }
94 
95       row_end = brows;
96       if (row_end > mbs) row_end = mbs;
97 
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     ierr = PetscFree(Jentry);CHKERRQ(ierr);
121     c->matentry = Jentry_new;
122   } else { /* ---------  c->htype == 'wp', use MatEntry2 ------------------*/
123     MatEntry2      *Jentry2_new,*Jentry2=c->matentry2;
124     ierr = PetscMalloc1(nz,&Jentry2_new);CHKERRQ(ierr);
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       }
130 
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       }
136 
137       row_end = brows;
138       if (row_end > mbs) row_end = mbs;
139 
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     ierr = PetscFree(Jentry2);CHKERRQ(ierr);
162     c->matentry2 = Jentry2_new;
163   } /* ---------------------------------------------*/
164 
165   ierr = PetscFree2(color_start,row_start);CHKERRQ(ierr);
166 
167   for (i=nbcols-1; i>0; i--) nrows_new[i] -= nrows_new[i-1];
168   ierr = PetscFree(c->nrows);CHKERRQ(ierr);
169   c->nrows = nrows_new;
170   PetscFunctionReturn(0);
171 }
172 
173 PetscErrorCode MatFDColoringSetUp_SeqXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
174 {
175   PetscErrorCode ierr;
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;
183 
184   PetscFunctionBegin;
185   ierr = ISColoringGetIS(iscoloring,PETSC_IGNORE,&isa);CHKERRQ(ierr);
186 
187   ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
188   ierr = PetscObjectTypeCompare((PetscObject)mat,MATSEQBAIJ,&isBAIJ);CHKERRQ(ierr);
189   ierr = PetscObjectTypeCompare((PetscObject)mat,MATSEQSELL,&isSELL);CHKERRQ(ierr);
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   }
205 
206   ierr       = PetscMalloc1(nis,&c->ncolumns);CHKERRQ(ierr);
207   ierr       = PetscMalloc1(nis,&c->columns);CHKERRQ(ierr);
208   ierr       = PetscMalloc1(nis,&c->nrows);CHKERRQ(ierr);
209   ierr       = PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));CHKERRQ(ierr);
210 
211   if (c->htype[0] == 'd') {
212     ierr       = PetscMalloc1(nz,&Jentry);CHKERRQ(ierr);
213     ierr       = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));CHKERRQ(ierr);
214     c->matentry = Jentry;
215   } else if (c->htype[0] == 'w') {
216     ierr       = PetscMalloc1(nz,&Jentry2);CHKERRQ(ierr);
217     ierr       = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry2));CHKERRQ(ierr);
218     c->matentry2 = Jentry2;
219   } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"htype is not supported");
220 
221   if (isBAIJ) {
222     ierr = MatGetColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
223   } else if (isSELL) {
224     ierr = MatGetColumnIJ_SeqSELL_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
225   } else {
226     ierr = MatGetColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
227   }
228 
229   ierr = PetscMalloc2(c->m,&rowhit,c->m,&valaddrhit);CHKERRQ(ierr);
230   ierr = PetscMemzero(rowhit,c->m*sizeof(PetscInt));CHKERRQ(ierr);
231 
232   nz = 0;
233   for (i=0; i<nis; i++) { /* loop over colors */
234     ierr = ISGetLocalSize(isa[i],&n);CHKERRQ(ierr);
235     ierr = ISGetIndices(isa[i],&is);CHKERRQ(ierr);
236 
237     c->ncolumns[i] = n;
238     if (n) {
239       ierr = PetscMalloc1(n,&c->columns[i]);CHKERRQ(ierr);
240       ierr = PetscLogObjectMemory((PetscObject)c,n*sizeof(PetscInt));CHKERRQ(ierr);
241       ierr = PetscMemcpy(c->columns[i],is,n*sizeof(PetscInt));CHKERRQ(ierr);
242     } else {
243       c->columns[i] = 0;
244     }
245 
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 */
260 
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     ierr = ISRestoreIndices(isa[i],&is);CHKERRQ(ierr);
282   }
283 
284   if (c->bcols > 1) {  /* reorder Jentry for faster MatFDColoringApply() */
285     ierr = MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz);CHKERRQ(ierr);
286   }
287 
288   if (isBAIJ) {
289     ierr = MatRestoreColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
290     ierr = PetscMalloc1(bs*mat->rmap->n,&c->dy);CHKERRQ(ierr);
291     ierr = PetscLogObjectMemory((PetscObject)c,bs*mat->rmap->n*sizeof(PetscScalar));CHKERRQ(ierr);
292   } else if (isSELL) {
293     ierr = MatRestoreColumnIJ_SeqSELL_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
294   } else {
295     ierr = MatRestoreColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
296   }
297   ierr = PetscFree2(rowhit,valaddrhit);CHKERRQ(ierr);
298   ierr = ISColoringRestoreIS(iscoloring,&isa);CHKERRQ(ierr);
299 
300   ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->rmap->n,PETSC_DETERMINE,0,NULL,&c->vscale);CHKERRQ(ierr);
301   ierr = PetscInfo3(c,"ncolors %D, brows %D and bcols %D are used.\n",c->ncolors,c->brows,c->bcols);CHKERRQ(ierr);
302   PetscFunctionReturn(0);
303 }
304