xref: /petsc/src/mat/impls/aij/seq/fdaij.c (revision bebe2cf65d55febe21a5af8db2bd2e168caaa2e7)
1 
2 #include <../src/mat/impls/aij/seq/aij.h>
3 #include <../src/mat/impls/baij/seq/baij.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 #undef __FUNCT__
11 #define __FUNCT__ "MatFDColoringCreate_SeqXAIJ"
12 PetscErrorCode MatFDColoringCreate_SeqXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
13 {
14   PetscErrorCode ierr;
15   PetscInt       bs,nis=iscoloring->n,m=mat->rmap->n;
16   PetscBool      isBAIJ;
17 
18   PetscFunctionBegin;
19   /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian */
20   ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
21   ierr = PetscObjectTypeCompare((PetscObject)mat,MATSEQBAIJ,&isBAIJ);CHKERRQ(ierr);
22   if (isBAIJ) {
23     c->brows = m;
24     c->bcols = 1;
25   } else { /* seqaij matrix */
26     /* bcols is chosen s.t. dy-array takes 50% of memory space as mat */
27     Mat_SeqAIJ *spA = (Mat_SeqAIJ*)mat->data;
28     PetscReal  mem;
29     PetscInt   nz,brows,bcols;
30 
31     bs    = 1; /* only bs=1 is supported for SeqAIJ matrix */
32 
33     nz    = spA->nz;
34     mem   = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
35     bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
36     brows = 1000/bcols;
37     if (bcols > nis) bcols = nis;
38     if (brows == 0 || brows > m) brows = m;
39     c->brows = brows;
40     c->bcols = bcols;
41   }
42 
43   c->M       = mat->rmap->N/bs;   /* set total rows, columns and local rows */
44   c->N       = mat->cmap->N/bs;
45   c->m       = mat->rmap->N/bs;
46   c->rstart  = 0;
47   c->ncolors = nis;
48   c->ctype   = IS_COLORING_GHOSTED;
49   PetscFunctionReturn(0);
50 }
51 
52 /*
53  Reorder Jentry such that blocked brows*bols of entries from dense matrix are inserted into Jacobian for improved cache performance
54    Input Parameters:
55 +  mat - the matrix containing the nonzero structure of the Jacobian
56 .  color - the coloring context
57 -  nz - number of local non-zeros in mat
58 */
59 #undef __FUNCT__
60 #define __FUNCT__ "MatFDColoringSetUpBlocked_AIJ_Private"
61 PetscErrorCode MatFDColoringSetUpBlocked_AIJ_Private(Mat mat,MatFDColoring c,PetscInt nz)
62 {
63   PetscErrorCode ierr;
64   PetscInt       i,j,nrows,nbcols,brows=c->brows,bcols=c->bcols,mbs=c->m,nis=c->ncolors;
65   PetscInt       *color_start,*row_start,*nrows_new,nz_new,row_end;
66 
67   PetscFunctionBegin;
68   if (brows < 1 || brows > mbs) brows = mbs;
69   ierr = PetscMalloc2(bcols+1,&color_start,bcols,&row_start);CHKERRQ(ierr);
70   ierr = PetscMalloc1(nis,&nrows_new);CHKERRQ(ierr);
71   ierr = PetscMalloc1(bcols*mat->rmap->n,&c->dy);CHKERRQ(ierr);
72   ierr = PetscLogObjectMemory((PetscObject)c,bcols*mat->rmap->n*sizeof(PetscScalar));CHKERRQ(ierr);
73 
74   nz_new = 0;
75   nbcols = 0;
76   color_start[bcols] = 0;
77 
78   if (c->htype[0] == 'd') { /* ----  c->htype == 'ds', use MatEntry --------*/
79     MatEntry       *Jentry_new,*Jentry=c->matentry;
80     ierr = PetscMalloc1(nz,&Jentry_new);CHKERRQ(ierr);
81     for (i=0; i<nis; i+=bcols) { /* loop over colors */
82       if (i + bcols > nis) {
83         color_start[nis - i] = color_start[bcols];
84         bcols                = nis - i;
85       }
86 
87       color_start[0] = color_start[bcols];
88       for (j=0; j<bcols; j++) {
89         color_start[j+1] = c->nrows[i+j] + color_start[j];
90         row_start[j]     = 0;
91       }
92 
93       row_end = brows;
94       if (row_end > mbs) row_end = mbs;
95 
96       while (row_end <= mbs) {   /* loop over block rows */
97         for (j=0; j<bcols; j++) {       /* loop over block columns */
98           nrows = c->nrows[i+j];
99           nz    = color_start[j];
100           while (row_start[j] < nrows) {
101             if (Jentry[nz].row >= row_end) {
102               color_start[j] = nz;
103               break;
104             } else { /* copy Jentry[nz] to Jentry_new[nz_new] */
105               Jentry_new[nz_new].row     = Jentry[nz].row + j*mbs; /* index in dy-array */
106               Jentry_new[nz_new].col     = Jentry[nz].col;
107               Jentry_new[nz_new].valaddr = Jentry[nz].valaddr;
108               nz_new++; nz++; row_start[j]++;
109             }
110           }
111         }
112         if (row_end == mbs) break;
113         row_end += brows;
114         if (row_end > mbs) row_end = mbs;
115       }
116       nrows_new[nbcols++] = nz_new;
117     }
118     ierr = PetscFree(Jentry);CHKERRQ(ierr);
119     c->matentry = Jentry_new;
120   } else { /* ---------  c->htype == 'wp', use MatEntry2 ------------------*/
121     MatEntry2      *Jentry2_new,*Jentry2=c->matentry2;
122     ierr = PetscMalloc1(nz,&Jentry2_new);CHKERRQ(ierr);
123     for (i=0; i<nis; i+=bcols) { /* loop over colors */
124       if (i + bcols > nis) {
125         color_start[nis - i] = color_start[bcols];
126         bcols                = nis - i;
127       }
128 
129       color_start[0] = color_start[bcols];
130       for (j=0; j<bcols; j++) {
131         color_start[j+1] = c->nrows[i+j] + color_start[j];
132         row_start[j]     = 0;
133       }
134 
135       row_end = brows;
136       if (row_end > mbs) row_end = mbs;
137 
138       while (row_end <= mbs) {   /* loop over block rows */
139         for (j=0; j<bcols; j++) {       /* loop over block columns */
140           nrows = c->nrows[i+j];
141           nz    = color_start[j];
142           while (row_start[j] < nrows) {
143             if (Jentry2[nz].row >= row_end) {
144               color_start[j] = nz;
145               break;
146             } else { /* copy Jentry2[nz] to Jentry2_new[nz_new] */
147               Jentry2_new[nz_new].row     = Jentry2[nz].row + j*mbs; /* index in dy-array */
148               Jentry2_new[nz_new].valaddr = Jentry2[nz].valaddr;
149               nz_new++; nz++; row_start[j]++;
150             }
151           }
152         }
153         if (row_end == mbs) break;
154         row_end += brows;
155         if (row_end > mbs) row_end = mbs;
156       }
157       nrows_new[nbcols++] = nz_new;
158     }
159     ierr = PetscFree(Jentry2);CHKERRQ(ierr);
160     c->matentry2 = Jentry2_new;
161   } /* ---------------------------------------------*/
162 
163   ierr = PetscFree2(color_start,row_start);CHKERRQ(ierr);
164 
165   for (i=nbcols-1; i>0; i--) nrows_new[i] -= nrows_new[i-1];
166   ierr = PetscFree(c->nrows);CHKERRQ(ierr);
167   c->nrows = nrows_new;
168   PetscFunctionReturn(0);
169 }
170 
171 #undef __FUNCT__
172 #define __FUNCT__ "MatFDColoringSetUp_SeqXAIJ"
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;
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   if (isBAIJ) {
190     Mat_SeqBAIJ *spA = (Mat_SeqBAIJ*)mat->data;
191     A_val = spA->a;
192     nz    = spA->nz;
193   } else {
194     Mat_SeqAIJ  *spA = (Mat_SeqAIJ*)mat->data;
195     A_val = spA->a;
196     nz    = spA->nz;
197     bs    = 1; /* only bs=1 is supported for SeqAIJ matrix */
198   }
199 
200   ierr       = PetscMalloc1(nis,&c->ncolumns);CHKERRQ(ierr);
201   ierr       = PetscMalloc1(nis,&c->columns);CHKERRQ(ierr);
202   ierr       = PetscMalloc1(nis,&c->nrows);CHKERRQ(ierr);
203   ierr       = PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));CHKERRQ(ierr);
204 
205   if (c->htype[0] == 'd') {
206     ierr       = PetscMalloc1(nz,&Jentry);CHKERRQ(ierr);
207     ierr       = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));CHKERRQ(ierr);
208     c->matentry = Jentry;
209   } else if (c->htype[0] == 'w') {
210     ierr       = PetscMalloc1(nz,&Jentry2);CHKERRQ(ierr);
211     ierr       = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry2));CHKERRQ(ierr);
212     c->matentry2 = Jentry2;
213   } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"htype is not supported");
214 
215   if (isBAIJ) {
216     ierr = MatGetColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
217   } else {
218     ierr = MatGetColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
219   }
220 
221   ierr = PetscMalloc2(c->m,&rowhit,c->m,&valaddrhit);CHKERRQ(ierr);
222   ierr = PetscMemzero(rowhit,c->m*sizeof(PetscInt));CHKERRQ(ierr);
223 
224   nz = 0;
225   for (i=0; i<nis; i++) { /* loop over colors */
226     ierr = ISGetLocalSize(isa[i],&n);CHKERRQ(ierr);
227     ierr = ISGetIndices(isa[i],&is);CHKERRQ(ierr);
228 
229     c->ncolumns[i] = n;
230     if (n) {
231       ierr = PetscMalloc1(n,&c->columns[i]);CHKERRQ(ierr);
232       ierr = PetscLogObjectMemory((PetscObject)c,n*sizeof(PetscInt));CHKERRQ(ierr);
233       ierr = PetscMemcpy(c->columns[i],is,n*sizeof(PetscInt));CHKERRQ(ierr);
234     } else {
235       c->columns[i] = 0;
236     }
237 
238     /* fast, crude version requires O(N*N) work */
239     bs2   = bs*bs;
240     nrows = 0;
241     for (j=0; j<n; j++) {  /* loop over columns */
242       col    = is[j];
243       row    = cj + ci[col];
244       m      = ci[col+1] - ci[col];
245       nrows += m;
246       for (k=0; k<m; k++) {  /* loop over columns marking them in rowhit */
247         rowhit[*row]       = col + 1;
248         valaddrhit[*row++] = &A_val[bs2*spidx[ci[col] + k]];
249       }
250     }
251     c->nrows[i] = nrows; /* total num of rows for this color */
252 
253     if (c->htype[0] == 'd') {
254       for (j=0; j<mbs; j++) { /* loop over rows */
255         if (rowhit[j]) {
256           Jentry[nz].row     = j;              /* local row index */
257           Jentry[nz].col     = rowhit[j] - 1;  /* local column index */
258           Jentry[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
259           nz++;
260           rowhit[j] = 0.0;                     /* zero rowhit for reuse */
261         }
262       }
263     }  else { /* c->htype == 'wp' */
264       for (j=0; j<mbs; j++) { /* loop over rows */
265         if (rowhit[j]) {
266           Jentry2[nz].row     = j;              /* local row index */
267           Jentry2[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
268           nz++;
269           rowhit[j] = 0.0;                     /* zero rowhit for reuse */
270         }
271       }
272     }
273     ierr = ISRestoreIndices(isa[i],&is);CHKERRQ(ierr);
274   }
275 
276   if (c->bcols > 1) {  /* reorder Jentry for faster MatFDColoringApply() */
277     ierr = MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz);CHKERRQ(ierr);
278   }
279 
280   if (isBAIJ) {
281     ierr = MatRestoreColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
282     ierr = PetscMalloc1(bs*mat->rmap->n,&c->dy);CHKERRQ(ierr);
283     ierr = PetscLogObjectMemory((PetscObject)c,bs*mat->rmap->n*sizeof(PetscScalar));CHKERRQ(ierr);
284   } else {
285     ierr = MatRestoreColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
286   }
287   ierr = PetscFree2(rowhit,valaddrhit);CHKERRQ(ierr);
288   ierr = ISColoringRestoreIS(iscoloring,&isa);CHKERRQ(ierr);
289 
290   ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->rmap->n,PETSC_DETERMINE,0,NULL,&c->vscale);CHKERRQ(ierr);
291   ierr = PetscInfo3(c,"ncolors %D, brows %D and bcols %D are used.\n",c->ncolors,c->brows,c->bcols);CHKERRQ(ierr);
292   PetscFunctionReturn(0);
293 }
294