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__ "MatFDColoringSetUp_SeqXAIJ" 11 PetscErrorCode MatFDColoringSetUp_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,nz_new,row_end,*row_start,*nrows_new; 21 PetscInt brows,bcols; 22 23 PetscFunctionBegin; 24 ierr = ISColoringGetIS(iscoloring,PETSC_IGNORE,&isa);CHKERRQ(ierr); 25 26 ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr); 27 ierr = PetscObjectTypeCompare((PetscObject)mat,MATSEQBAIJ,&isBAIJ);CHKERRQ(ierr); 28 if (isBAIJ) { 29 Mat_SeqBAIJ *spA = (Mat_SeqBAIJ*)mat->data; 30 A_val = spA->a; 31 nz = spA->nz; 32 } else { 33 Mat_SeqAIJ *spA = (Mat_SeqAIJ*)mat->data; 34 A_val = spA->a; 35 nz = spA->nz; 36 bs = 1; /* only bs=1 is supported for SeqAIJ matrix */ 37 } 38 39 N = mat->cmap->N/bs; 40 c->M = mat->rmap->N/bs; /* set total rows, columns and local rows */ 41 c->N = mat->cmap->N/bs; 42 c->m = mat->rmap->N/bs; 43 c->rstart = 0; 44 45 c->ncolors = nis; 46 ierr = PetscMalloc(nis*sizeof(PetscInt),&c->ncolumns);CHKERRQ(ierr); 47 ierr = PetscMalloc(nis*sizeof(PetscInt*),&c->columns);CHKERRQ(ierr); 48 ierr = PetscMalloc(nis*sizeof(PetscInt),&c->nrows);CHKERRQ(ierr); 49 ierr = PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));CHKERRQ(ierr); 50 51 ierr = PetscMalloc(nz*sizeof(MatEntry),&Jentry);CHKERRQ(ierr); 52 ierr = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));CHKERRQ(ierr); 53 c->matentry = Jentry; 54 55 ierr = PetscMalloc2(nis+1,PetscInt,&color_start,nis+1,PetscInt,&row_start);CHKERRQ(ierr); 56 57 if (isBAIJ) { 58 ierr = MatGetColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr); 59 } else { 60 ierr = MatGetColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr); 61 } 62 63 ierr = PetscMalloc2(c->m,PetscInt,&rowhit,c->m,PetscScalar*,&valaddrhit);CHKERRQ(ierr); 64 ierr = PetscMemzero(rowhit,c->m*sizeof(PetscInt));CHKERRQ(ierr); 65 66 // estimate bcol 67 PetscReal mem; 68 mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*c->m*sizeof(PetscInt); 69 bcols = (PetscInt)(0.5*mem /(c->m*sizeof(PetscScalar))); 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 if (!isBAIJ && bcols > 1) { 119 PetscInt nbcols = 0; 120 121 ierr = PetscMalloc(nz*sizeof(MatEntry),&Jentry_new);CHKERRQ(ierr); 122 for (i=0; i<nis; i++) row_start[i] = 0; 123 124 125 ierr = PetscOptionsInt("-bcols","The number of block columns","",bcols,&bcols,NULL);CHKERRQ(ierr); 126 if (bcols > nis) bcols = nis; 127 c->brows = brows; 128 c->bcols = bcols; 129 130 ierr = PetscMalloc(nis*sizeof(PetscInt),&nrows_new);CHKERRQ(ierr); 131 132 nz_new = 0; 133 for (i=0; i<nis; i+=bcols) { /* loop over colors */ 134 if (i + bcols > nis) bcols = nis - i; 135 136 row_end = brows; 137 m = mat->rmap->n; 138 if (row_end > m) row_end = m; 139 while (row_end <= m) { /* loop over block rows */ 140 for (j=0; j<bcols; j++) { /* loop over block columns */ 141 nrows = c->nrows[i+j]; 142 for (nz=color_start[i+j]; nz<color_start[i+j+1]; nz++) { /* for each Jentry */ 143 if (row_start[i+j] >= nrows) break; 144 if (Jentry[nz].row >= row_end) { 145 color_start[i+j] = nz; 146 break; 147 } else { 148 Jentry_new[nz_new].row = Jentry[nz].row + j*m; /* index in dy-array */ 149 Jentry_new[nz_new].col = Jentry[nz].col; 150 Jentry_new[nz_new].valaddr = Jentry[nz].valaddr; 151 nz_new++; 152 row_start[i+j]++; 153 } 154 } 155 } 156 if (row_end == m) break; 157 row_end += brows; 158 if (row_end > m) row_end = m; 159 } 160 nrows_new[nbcols++] = nz_new; 161 } 162 for (i=nbcols-1; i>0; i--) nrows_new[i] -= nrows_new[i-1]; 163 ierr = PetscFree(c->nrows);CHKERRQ(ierr); 164 c->nrows = nrows_new; 165 166 ierr = PetscFree(Jentry);CHKERRQ(ierr); 167 c->matentry = Jentry_new; 168 ierr = PetscMalloc(c->bcols*mat->rmap->n*sizeof(PetscScalar),&c->dy);CHKERRQ(ierr); 169 } 170 //--------------------------------------- 171 172 if (isBAIJ) { 173 ierr = MatRestoreColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr); 174 ierr = PetscMalloc(bs*mat->rmap->n*sizeof(PetscScalar),&c->dy);CHKERRQ(ierr); 175 } else { 176 ierr = MatRestoreColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr); 177 } 178 ierr = PetscFree2(rowhit,valaddrhit);CHKERRQ(ierr); 179 ierr = ISColoringRestoreIS(iscoloring,&isa);CHKERRQ(ierr); 180 ierr = PetscFree2(color_start,row_start);CHKERRQ(ierr); 181 182 c->ctype = IS_COLORING_GHOSTED; 183 ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->rmap->n,PETSC_DETERMINE,0,NULL,&c->vscale);CHKERRQ(ierr); 184 PetscFunctionReturn(0); 185 } 186