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