xref: /petsc/src/mat/impls/aij/seq/fdaij.c (revision 93dfae198b7b88e3d1d0b167d98d82ce5d6af8d8)
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,nz,bcols,nis=iscoloring->n;
15   PetscBool      isBAIJ;
16   PetscReal      mem;
17 
18   PetscFunctionBegin;
19   ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
20   ierr = PetscObjectTypeCompare((PetscObject)mat,MATSEQBAIJ,&isBAIJ);CHKERRQ(ierr);
21   if (isBAIJ) {
22     Mat_SeqBAIJ *spA = (Mat_SeqBAIJ*)mat->data;
23     nz    = spA->nz;
24   } else {
25     Mat_SeqAIJ  *spA = (Mat_SeqAIJ*)mat->data;
26     nz    = spA->nz;
27     bs    = 1; /* only bs=1 is supported for SeqAIJ matrix */
28   }
29   c->M       = mat->rmap->N/bs;   /* set total rows, columns and local rows */
30   c->N       = mat->cmap->N/bs;
31   c->m       = mat->rmap->N/bs;
32   c->rstart  = 0;
33   c->ncolors = nis;
34 
35   /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian;
36      bcols is chosen s.t. dy-array takes 50% of memory space as mat */
37   mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*c->m*sizeof(PetscInt);
38   bcols = (PetscInt)(0.5*mem /(c->m*sizeof(PetscScalar)));
39   if (bcols > nis) bcols = nis;
40   c->brows = 1000/bcols;
41   c->bcols = bcols;
42   c->ctype = IS_COLORING_GHOSTED;
43   PetscFunctionReturn(0);
44 }
45 
46 #undef __FUNCT__
47 #define __FUNCT__ "MatFDColoringSetUp_SeqXAIJ"
48 PetscErrorCode MatFDColoringSetUp_SeqXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
49 {
50   PetscErrorCode ierr;
51   PetscInt       i,n,nrows,N,j,k,m,ncols,col,nis=iscoloring->n,*rowhit,bs,bs2,*spidx,nz;
52   const PetscInt *is,*row,*ci,*cj;
53   IS             *isa;
54   PetscBool      isBAIJ;
55   PetscScalar    *A_val,**valaddrhit;
56   MatEntry       *Jentry,*Jentry_new;
57   PetscInt       *color_start,nz_new,row_end,*row_start,*nrows_new;
58   PetscInt       bcols=c->bcols;
59 
60   PetscFunctionBegin;
61   ierr = ISColoringGetIS(iscoloring,PETSC_IGNORE,&isa);CHKERRQ(ierr);
62 
63   ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
64   ierr = PetscObjectTypeCompare((PetscObject)mat,MATSEQBAIJ,&isBAIJ);CHKERRQ(ierr);
65   if (isBAIJ) {
66     Mat_SeqBAIJ *spA = (Mat_SeqBAIJ*)mat->data;
67     A_val = spA->a;
68     nz    = spA->nz;
69   } else {
70     Mat_SeqAIJ  *spA = (Mat_SeqAIJ*)mat->data;
71     A_val = spA->a;
72     nz    = spA->nz;
73     bs    = 1; /* only bs=1 is supported for SeqAIJ matrix */
74   }
75 
76   N         = mat->cmap->N/bs;
77   ierr       = PetscMalloc(nis*sizeof(PetscInt),&c->ncolumns);CHKERRQ(ierr);
78   ierr       = PetscMalloc(nis*sizeof(PetscInt*),&c->columns);CHKERRQ(ierr);
79   ierr       = PetscMalloc(nis*sizeof(PetscInt),&c->nrows);CHKERRQ(ierr);
80   ierr       = PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));CHKERRQ(ierr);
81 
82   ierr       = PetscMalloc(nz*sizeof(MatEntry),&Jentry);CHKERRQ(ierr);
83   ierr       = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));CHKERRQ(ierr);
84   c->matentry = Jentry;
85 
86   if (isBAIJ) {
87     ierr = MatGetColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
88   } else {
89     ierr = MatGetColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
90   }
91 
92   ierr = PetscMalloc3(c->m,PetscInt,&rowhit,c->m,PetscScalar*,&valaddrhit,nis+1,PetscInt,&color_start);CHKERRQ(ierr);
93   ierr = PetscMemzero(rowhit,c->m*sizeof(PetscInt));CHKERRQ(ierr);
94 
95   nz = 0;
96   for (i=0; i<nis; i++) { /* loop over colors */
97     color_start[i] = nz;
98 
99     ierr = ISGetLocalSize(isa[i],&n);CHKERRQ(ierr);
100     ierr = ISGetIndices(isa[i],&is);CHKERRQ(ierr);
101 
102     c->ncolumns[i] = n;
103     if (n) {
104       ierr = PetscMalloc(n*sizeof(PetscInt),&c->columns[i]);CHKERRQ(ierr);
105       ierr = PetscLogObjectMemory((PetscObject)c,n*sizeof(PetscInt));CHKERRQ(ierr);
106       ierr = PetscMemcpy(c->columns[i],is,n*sizeof(PetscInt));CHKERRQ(ierr);
107     } else {
108       c->columns[i] = 0;
109     }
110 
111     /* fast, crude version requires O(N*N) work */
112     bs2   = bs*bs;
113     nrows = 0;
114     for (j=0; j<n; j++) {  /* loop over columns */
115       col    = is[j];
116       row    = cj + ci[col];
117       m      = ci[col+1] - ci[col];
118       nrows += m;
119       for (k=0; k<m; k++) {  /* loop over columns marking them in rowhit */
120         rowhit[*row]       = col + 1;
121         valaddrhit[*row++] = &A_val[bs2*spidx[ci[col] + k]];
122       }
123     }
124     c->nrows[i] = nrows; /* total num of rows for this color */
125 
126     for (j=0; j<N; j++) { /* loop over rows */
127       if (rowhit[j]) {
128         Jentry[nz].row     = j;              /* local row index */
129         Jentry[nz].col     = rowhit[j] - 1;  /* local column index */
130         Jentry[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
131         nz++;
132         rowhit[j] = 0.0;                     /* zero rowhit for reuse */
133       }
134     }
135     ierr = ISRestoreIndices(isa[i],&is);CHKERRQ(ierr);
136   }
137   color_start[nis] = nz;
138 
139   // ---------- reorder Jentry ------------
140   if (!isBAIJ && bcols > 1) {
141     PetscInt nbcols=0,brows=c->brows;
142 
143     m = mat->rmap->n;
144     if (brows < 1) brows = m;
145 
146     ierr = PetscMalloc(nz*sizeof(MatEntry),&Jentry_new);CHKERRQ(ierr);
147     ierr = PetscMalloc(bcols*sizeof(PetscInt),&row_start);CHKERRQ(ierr);
148     ierr = PetscMalloc(nis*sizeof(PetscInt),&nrows_new);CHKERRQ(ierr);
149 
150     nz_new  = 0;
151     for (i=0; i<nis; i+=bcols) { /* loop over colors */
152       if (i + bcols > nis) bcols = nis - i;
153 
154       row_end = brows;
155       if (row_end > m) row_end = m;
156       for (j=0; j<bcols; j++) row_start[j] = 0;
157       while (row_end <= m) { /* loop over block rows */
158         for (j=0; j<bcols; j++) {       /* loop over block columns */
159           nrows = c->nrows[i+j];
160           for (nz=color_start[i+j]; nz<color_start[i+j+1]; nz++) { /* for each Jentry */
161             if (row_start[j] >= nrows) break;
162             if (Jentry[nz].row >= row_end) {
163               color_start[i+j] = nz;
164               break;
165             } else {
166               Jentry_new[nz_new].row     = Jentry[nz].row + j*m; /* index in dy-array */
167               Jentry_new[nz_new].col     = Jentry[nz].col;
168               Jentry_new[nz_new].valaddr = Jentry[nz].valaddr;
169               nz_new++;
170               row_start[j]++;
171             }
172           }
173         }
174         if (row_end == m) break;
175         row_end += brows;
176         if (row_end > m) row_end = m;
177       }
178       nrows_new[nbcols++] = nz_new;
179     }
180     for (i=nbcols-1; i>0; i--) nrows_new[i] -= nrows_new[i-1];
181     ierr = PetscFree(c->nrows);CHKERRQ(ierr);
182     c->nrows = nrows_new;
183 
184     ierr = PetscFree(Jentry);CHKERRQ(ierr);
185     c->matentry = Jentry_new;
186     ierr = PetscMalloc(c->bcols*mat->rmap->n*sizeof(PetscScalar),&c->dy);CHKERRQ(ierr);
187     ierr = PetscFree(row_start);CHKERRQ(ierr);
188   }
189   //---------------------------------------
190 
191   if (isBAIJ) {
192     ierr = MatRestoreColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
193     ierr = PetscMalloc(bs*mat->rmap->n*sizeof(PetscScalar),&c->dy);CHKERRQ(ierr);
194   } else {
195     ierr = MatRestoreColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr);
196   }
197   ierr = PetscFree3(rowhit,valaddrhit,color_start);CHKERRQ(ierr);
198   ierr = ISColoringRestoreIS(iscoloring,&isa);CHKERRQ(ierr);
199 
200   ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->rmap->n,PETSC_DETERMINE,0,NULL,&c->vscale);CHKERRQ(ierr);
201   PetscFunctionReturn(0);
202 }
203