xref: /petsc/src/mat/impls/aij/seq/fdaij.c (revision 054951ad227a3627c7f2e3d3160721959c51d9d8)
1 
2 #include <../src/mat/impls/aij/seq/aij.h>
3 #include <../src/mat/impls/baij/seq/baij.h>
4 
5 /*
6     This routine is shared by SeqAIJ and SeqBAIJ matrices,
7     since it operators only on the nonzero structure of the elements or blocks.
8 */
9 #undef __FUNCT__
10 #define __FUNCT__ "MatFDColoringCreate_SeqXAIJ"
11 PetscErrorCode MatFDColoringCreate_SeqXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
12 {
13   PetscErrorCode ierr;
14   PetscInt       bs,nis=iscoloring->n,m=mat->rmap->n;
15   PetscBool      isBAIJ;
16 
17   PetscFunctionBegin;
18   /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian */
19   ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
20   ierr = PetscObjectTypeCompare((PetscObject)mat,MATSEQBAIJ,&isBAIJ);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     Mat_SeqAIJ *spA = (Mat_SeqAIJ*)mat->data;
27     PetscReal  mem;
28     PetscInt   nz,brows,bcols;
29 
30     bs    = 1; /* only bs=1 is supported for SeqAIJ matrix */
31 
32     nz    = spA->nz;
33     mem   = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
34     bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
35     brows = 1000/bcols;
36     if (bcols > nis) bcols = nis;
37     if (brows == 0 || brows > m) brows = m;
38     c->brows = brows;
39     c->bcols = bcols;
40   }
41 
42   c->M       = mat->rmap->N/bs;   /* set total rows, columns and local rows */
43   c->N       = mat->cmap->N/bs;
44   c->m       = mat->rmap->N/bs;
45   c->rstart  = 0;
46   c->ncolors = nis;
47   c->ctype   = IS_COLORING_GHOSTED;
48   PetscFunctionReturn(0);
49 }
50 
51 /*
52  Reorder Jentry such that blocked brows*bols of entries from dense matrix are inserted into sparse Jacobian
53    Input Parameters:
54 +  mat - the matrix containing the nonzero structure of the Jacobian
55 .  color - the coloring context
56 -  nz - number of local non-zeros in mat
57 */
58 #undef __FUNCT__
59 #define __FUNCT__ "MatFDColoringSetUpBlocked_AIJ_Private"
60 PetscErrorCode MatFDColoringSetUpBlocked_AIJ_Private(Mat mat,MatFDColoring c,PetscInt nz)
61 {
62   PetscErrorCode ierr;
63   PetscInt       i,j,nrows,nbcols,brows=c->brows,bcols=c->bcols,mbs=c->m,nis=c->ncolors;
64   PetscInt       *color_start,*row_start,*nrows_new,nz_new,row_end;
65   MatEntry       *Jentry_new,*Jentry=c->matentry;
66 
67   PetscFunctionBegin;
68   if (brows < 1 || brows > mbs) brows = mbs;
69   ierr = PetscMalloc2(bcols+1,PetscInt,&color_start,bcols,PetscInt,&row_start);CHKERRQ(ierr);
70   ierr = PetscMalloc(nz*sizeof(MatEntry),&Jentry_new);CHKERRQ(ierr);
71   ierr = PetscMalloc(nis*sizeof(PetscInt),&nrows_new);CHKERRQ(ierr);
72   ierr = PetscMalloc(bcols*mat->rmap->n*sizeof(PetscScalar),&c->dy);CHKERRQ(ierr);
73   ierr = PetscLogObjectMemory((PetscObject)c,bcols*mat->rmap->n*sizeof(PetscScalar));CHKERRQ(ierr);
74 
75   nz_new = 0;
76   nbcols = 0;
77   color_start[bcols] = 0;
78   for (i=0; i<nis; i+=bcols) { /* loop over colors */
79     if (i + bcols > nis) {
80       color_start[nis - i] = color_start[bcols];
81       bcols                = nis - i;
82     }
83 
84     color_start[0] = color_start[bcols];
85     for (j=0; j<bcols; j++) {
86       color_start[j+1] = c->nrows[i+j] + color_start[j];
87       row_start[j]     = 0;
88     }
89 
90     row_end = brows;
91     if (row_end > mbs) row_end = mbs;
92 
93     while (row_end <= mbs) {   /* loop over block rows */
94       for (j=0; j<bcols; j++) {       /* loop over block columns */
95         nrows = c->nrows[i+j];
96         nz    = color_start[j];
97         while (row_start[j] < nrows) {
98           if (Jentry[nz].row >= row_end) {
99             color_start[j] = nz;
100             break;
101           } else { /* copy Jentry[nz] to Jentry_new[nz_new] */
102             Jentry_new[nz_new].row     = Jentry[nz].row + j*mbs; /* index in dy-array */
103             Jentry_new[nz_new].col     = Jentry[nz].col;
104             Jentry_new[nz_new].valaddr = Jentry[nz].valaddr;
105             nz_new++; nz++; row_start[j]++;
106           }
107         }
108       }
109       if (row_end == mbs) break;
110       row_end += brows;
111       if (row_end > mbs) row_end = mbs;
112     }
113     nrows_new[nbcols++] = nz_new;
114   }
115   ierr = PetscFree2(color_start,row_start);CHKERRQ(ierr);
116 
117   for (i=nbcols-1; i>0; i--) nrows_new[i] -= nrows_new[i-1];
118   ierr = PetscFree(c->nrows);CHKERRQ(ierr);
119   ierr = PetscFree(Jentry);CHKERRQ(ierr);
120   c->nrows    = nrows_new;
121   c->matentry = Jentry_new;
122   PetscFunctionReturn(0);
123 }
124 
125 #undef __FUNCT__
126 #define __FUNCT__ "MatFDColoringSetUp_SeqXAIJ"
127 PetscErrorCode MatFDColoringSetUp_SeqXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
128 {
129   PetscErrorCode ierr;
130   PetscInt       i,n,nrows,mbs=c->m,j,k,m,ncols,col,nis=iscoloring->n,*rowhit,bs,bs2,*spidx,nz;
131   const PetscInt *is,*row,*ci,*cj;
132   IS             *isa;
133   PetscBool      isBAIJ;
134   PetscScalar    *A_val,**valaddrhit;
135   MatEntry       *Jentry;
136 
137   PetscFunctionBegin;
138   ierr = ISColoringGetIS(iscoloring,PETSC_IGNORE,&isa);CHKERRQ(ierr);
139 
140   ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
141   ierr = PetscObjectTypeCompare((PetscObject)mat,MATSEQBAIJ,&isBAIJ);CHKERRQ(ierr);
142   if (isBAIJ) {
143     Mat_SeqBAIJ *spA = (Mat_SeqBAIJ*)mat->data;
144     A_val = spA->a;
145     nz    = spA->nz;
146   } else {
147     Mat_SeqAIJ  *spA = (Mat_SeqAIJ*)mat->data;
148     A_val = spA->a;
149     nz    = spA->nz;
150     bs    = 1; /* only bs=1 is supported for SeqAIJ matrix */
151   }
152 
153   ierr       = PetscMalloc(nis*sizeof(PetscInt),&c->ncolumns);CHKERRQ(ierr);
154   ierr       = PetscMalloc(nis*sizeof(PetscInt*),&c->columns);CHKERRQ(ierr);
155   ierr       = PetscMalloc(nis*sizeof(PetscInt),&c->nrows);CHKERRQ(ierr);
156   ierr       = PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));CHKERRQ(ierr);
157 
158   ierr       = PetscMalloc(nz*sizeof(MatEntry),&Jentry);CHKERRQ(ierr);
159   ierr       = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));CHKERRQ(ierr);
160   c->matentry = Jentry;
161 
162   if (isBAIJ) {
163     ierr = MatGetColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
164   } else {
165     ierr = MatGetColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
166   }
167 
168   ierr = PetscMalloc2(c->m,PetscInt,&rowhit,c->m,PetscScalar*,&valaddrhit);CHKERRQ(ierr);
169   ierr = PetscMemzero(rowhit,c->m*sizeof(PetscInt));CHKERRQ(ierr);
170 
171   nz = 0;
172   for (i=0; i<nis; i++) { /* loop over colors */
173     ierr = ISGetLocalSize(isa[i],&n);CHKERRQ(ierr);
174     ierr = ISGetIndices(isa[i],&is);CHKERRQ(ierr);
175 
176     c->ncolumns[i] = n;
177     if (n) {
178       ierr = PetscMalloc(n*sizeof(PetscInt),&c->columns[i]);CHKERRQ(ierr);
179       ierr = PetscLogObjectMemory((PetscObject)c,n*sizeof(PetscInt));CHKERRQ(ierr);
180       ierr = PetscMemcpy(c->columns[i],is,n*sizeof(PetscInt));CHKERRQ(ierr);
181     } else {
182       c->columns[i] = 0;
183     }
184 
185     /* fast, crude version requires O(N*N) work */
186     bs2   = bs*bs;
187     nrows = 0;
188     for (j=0; j<n; j++) {  /* loop over columns */
189       col    = is[j];
190       row    = cj + ci[col];
191       m      = ci[col+1] - ci[col];
192       nrows += m;
193       for (k=0; k<m; k++) {  /* loop over columns marking them in rowhit */
194         rowhit[*row]       = col + 1;
195         valaddrhit[*row++] = &A_val[bs2*spidx[ci[col] + k]];
196       }
197     }
198     c->nrows[i] = nrows; /* total num of rows for this color */
199 
200     for (j=0; j<mbs; j++) { /* loop over rows */
201       if (rowhit[j]) {
202         Jentry[nz].row     = j;              /* local row index */
203         Jentry[nz].col     = rowhit[j] - 1;  /* local column index */
204         Jentry[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
205         nz++;
206         rowhit[j] = 0.0;                     /* zero rowhit for reuse */
207       }
208     }
209     ierr = ISRestoreIndices(isa[i],&is);CHKERRQ(ierr);
210   }
211 
212   if (c->bcols > 1) {  /* reorder Jentry for faster MatFDColoringApply() */
213     ierr = MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz);CHKERRQ(ierr);
214   }
215 
216   if (isBAIJ) {
217     ierr = MatRestoreColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
218     ierr = PetscMalloc(bs*mat->rmap->n*sizeof(PetscScalar),&c->dy);CHKERRQ(ierr);
219     ierr = PetscLogObjectMemory((PetscObject)c,bs*mat->rmap->n*sizeof(PetscScalar));CHKERRQ(ierr);
220   } else {
221     ierr = MatRestoreColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
222   }
223   ierr = PetscFree2(rowhit,valaddrhit);CHKERRQ(ierr);
224   ierr = ISColoringRestoreIS(iscoloring,&isa);CHKERRQ(ierr);
225 
226   ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->rmap->n,PETSC_DETERMINE,0,NULL,&c->vscale);CHKERRQ(ierr);
227 #if defined(PETSC_USE_INFO)
228   ierr = PetscInfo3(c,"ncolors %D, brows %D and bcols %D are used.\n",c->ncolors,c->brows,c->bcols);CHKERRQ(ierr);
229 #endif
230   PetscFunctionReturn(0);
231 }
232