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