1 #include <../src/mat/impls/aij/seq/aij.h> 2 #include <../src/mat/impls/baij/seq/baij.h> 3 #include <../src/mat/impls/sell/seq/sell.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 PetscErrorCode MatFDColoringCreate_SeqXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c) 11 { 12 PetscErrorCode ierr; 13 PetscInt bs,nis=iscoloring->n,m=mat->rmap->n; 14 PetscBool isBAIJ,isSELL; 15 16 PetscFunctionBegin; 17 /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian */ 18 ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr); 19 ierr = PetscObjectBaseTypeCompare((PetscObject)mat,MATSEQBAIJ,&isBAIJ);CHKERRQ(ierr); 20 ierr = PetscObjectTypeCompare((PetscObject)mat,MATSEQSELL,&isSELL);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 PetscReal mem; 27 PetscInt nz,brows,bcols; 28 if (isSELL) { 29 Mat_SeqSELL *spA = (Mat_SeqSELL*)mat->data; 30 nz = spA->nz; 31 } else { 32 Mat_SeqAIJ *spA = (Mat_SeqAIJ*)mat->data; 33 nz = spA->nz; 34 } 35 36 bs = 1; /* only bs=1 is supported for SeqAIJ matrix */ 37 mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt); 38 bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar))); 39 brows = 1000/bcols; 40 if (bcols > nis) bcols = nis; 41 if (brows == 0 || brows > m) brows = m; 42 c->brows = brows; 43 c->bcols = bcols; 44 } 45 46 c->M = mat->rmap->N/bs; /* set total rows, columns and local rows */ 47 c->N = mat->cmap->N/bs; 48 c->m = mat->rmap->N/bs; 49 c->rstart = 0; 50 c->ncolors = nis; 51 c->ctype = iscoloring->ctype; 52 PetscFunctionReturn(0); 53 } 54 55 /* 56 Reorder Jentry such that blocked brows*bols of entries from dense matrix are inserted into Jacobian for improved cache performance 57 Input Parameters: 58 + mat - the matrix containing the nonzero structure of the Jacobian 59 . color - the coloring context 60 - nz - number of local non-zeros in mat 61 */ 62 PetscErrorCode MatFDColoringSetUpBlocked_AIJ_Private(Mat mat,MatFDColoring c,PetscInt nz) 63 { 64 PetscErrorCode ierr; 65 PetscInt i,j,nrows,nbcols,brows=c->brows,bcols=c->bcols,mbs=c->m,nis=c->ncolors; 66 PetscInt *color_start,*row_start,*nrows_new,nz_new,row_end; 67 68 PetscFunctionBegin; 69 if (brows < 1 || brows > mbs) brows = mbs; 70 ierr = PetscMalloc2(bcols+1,&color_start,bcols,&row_start);CHKERRQ(ierr); 71 ierr = PetscCalloc1(nis,&nrows_new);CHKERRQ(ierr); 72 ierr = PetscMalloc1(bcols*mat->rmap->n,&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 79 if (c->htype[0] == 'd') { /* ---- c->htype == 'ds', use MatEntry --------*/ 80 MatEntry *Jentry_new,*Jentry=c->matentry; 81 82 ierr = PetscMalloc1(nz,&Jentry_new);CHKERRQ(ierr); 83 for (i=0; i<nis; i+=bcols) { /* loop over colors */ 84 if (i + bcols > nis) { 85 color_start[nis - i] = color_start[bcols]; 86 bcols = nis - i; 87 } 88 89 color_start[0] = color_start[bcols]; 90 for (j=0; j<bcols; j++) { 91 color_start[j+1] = c->nrows[i+j] + color_start[j]; 92 row_start[j] = 0; 93 } 94 95 row_end = brows; 96 if (row_end > mbs) row_end = mbs; 97 98 while (row_end <= mbs) { /* loop over block rows */ 99 for (j=0; j<bcols; j++) { /* loop over block columns */ 100 nrows = c->nrows[i+j]; 101 nz = color_start[j]; 102 while (row_start[j] < nrows) { 103 if (Jentry[nz].row >= row_end) { 104 color_start[j] = nz; 105 break; 106 } else { /* copy Jentry[nz] to Jentry_new[nz_new] */ 107 Jentry_new[nz_new].row = Jentry[nz].row + j*mbs; /* index in dy-array */ 108 Jentry_new[nz_new].col = Jentry[nz].col; 109 Jentry_new[nz_new].valaddr = Jentry[nz].valaddr; 110 nz_new++; nz++; row_start[j]++; 111 } 112 } 113 } 114 if (row_end == mbs) break; 115 row_end += brows; 116 if (row_end > mbs) row_end = mbs; 117 } 118 nrows_new[nbcols++] = nz_new; 119 } 120 ierr = PetscFree(Jentry);CHKERRQ(ierr); 121 c->matentry = Jentry_new; 122 } else { /* --------- c->htype == 'wp', use MatEntry2 ------------------*/ 123 MatEntry2 *Jentry2_new,*Jentry2=c->matentry2; 124 125 ierr = PetscMalloc1(nz,&Jentry2_new);CHKERRQ(ierr); 126 for (i=0; i<nis; i+=bcols) { /* loop over colors */ 127 if (i + bcols > nis) { 128 color_start[nis - i] = color_start[bcols]; 129 bcols = nis - i; 130 } 131 132 color_start[0] = color_start[bcols]; 133 for (j=0; j<bcols; j++) { 134 color_start[j+1] = c->nrows[i+j] + color_start[j]; 135 row_start[j] = 0; 136 } 137 138 row_end = brows; 139 if (row_end > mbs) row_end = mbs; 140 141 while (row_end <= mbs) { /* loop over block rows */ 142 for (j=0; j<bcols; j++) { /* loop over block columns */ 143 nrows = c->nrows[i+j]; 144 nz = color_start[j]; 145 while (row_start[j] < nrows) { 146 if (Jentry2[nz].row >= row_end) { 147 color_start[j] = nz; 148 break; 149 } else { /* copy Jentry2[nz] to Jentry2_new[nz_new] */ 150 Jentry2_new[nz_new].row = Jentry2[nz].row + j*mbs; /* index in dy-array */ 151 Jentry2_new[nz_new].valaddr = Jentry2[nz].valaddr; 152 nz_new++; nz++; row_start[j]++; 153 } 154 } 155 } 156 if (row_end == mbs) break; 157 row_end += brows; 158 if (row_end > mbs) row_end = mbs; 159 } 160 nrows_new[nbcols++] = nz_new; 161 } 162 ierr = PetscFree(Jentry2);CHKERRQ(ierr); 163 c->matentry2 = Jentry2_new; 164 } /* ---------------------------------------------*/ 165 166 ierr = PetscFree2(color_start,row_start);CHKERRQ(ierr); 167 168 for (i=nbcols-1; i>0; i--) nrows_new[i] -= nrows_new[i-1]; 169 ierr = PetscFree(c->nrows);CHKERRQ(ierr); 170 c->nrows = nrows_new; 171 PetscFunctionReturn(0); 172 } 173 174 PetscErrorCode MatFDColoringSetUp_SeqXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c) 175 { 176 PetscErrorCode ierr; 177 PetscInt i,n,nrows,mbs=c->m,j,k,m,ncols,col,nis=iscoloring->n,*rowhit,bs,bs2,*spidx,nz,tmp; 178 const PetscInt *is,*row,*ci,*cj; 179 PetscBool isBAIJ,isSELL; 180 const PetscScalar *A_val; 181 PetscScalar **valaddrhit; 182 MatEntry *Jentry; 183 MatEntry2 *Jentry2; 184 185 PetscFunctionBegin; 186 ierr = ISColoringGetIS(iscoloring,PETSC_OWN_POINTER,PETSC_IGNORE,&c->isa);CHKERRQ(ierr); 187 188 ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr); 189 ierr = PetscObjectBaseTypeCompare((PetscObject)mat,MATSEQBAIJ,&isBAIJ);CHKERRQ(ierr); 190 ierr = PetscObjectTypeCompare((PetscObject)mat,MATSEQSELL,&isSELL);CHKERRQ(ierr); 191 if (isBAIJ) { 192 Mat_SeqBAIJ *spA = (Mat_SeqBAIJ*)mat->data; 193 194 A_val = spA->a; 195 nz = spA->nz; 196 } else if (isSELL) { 197 Mat_SeqSELL *spA = (Mat_SeqSELL*)mat->data; 198 199 A_val = spA->val; 200 nz = spA->nz; 201 bs = 1; /* only bs=1 is supported for SeqSELL matrix */ 202 } else { 203 Mat_SeqAIJ *spA = (Mat_SeqAIJ*)mat->data; 204 205 A_val = spA->a; 206 nz = spA->nz; 207 bs = 1; /* only bs=1 is supported for SeqAIJ matrix */ 208 } 209 210 ierr = PetscMalloc2(nis,&c->ncolumns,nis,&c->columns);CHKERRQ(ierr); 211 ierr = PetscMalloc1(nis,&c->nrows);CHKERRQ(ierr); /* nrows is freeed separately from ncolumns and columns */ 212 ierr = PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));CHKERRQ(ierr); 213 214 if (c->htype[0] == 'd') { 215 ierr = PetscMalloc1(nz,&Jentry);CHKERRQ(ierr); 216 ierr = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));CHKERRQ(ierr); 217 c->matentry = Jentry; 218 } else if (c->htype[0] == 'w') { 219 ierr = PetscMalloc1(nz,&Jentry2);CHKERRQ(ierr); 220 ierr = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry2));CHKERRQ(ierr); 221 c->matentry2 = Jentry2; 222 } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"htype is not supported"); 223 224 if (isBAIJ) { 225 ierr = MatGetColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr); 226 } else if (isSELL) { 227 ierr = MatGetColumnIJ_SeqSELL_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr); 228 } else { 229 ierr = MatGetColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr); 230 } 231 232 ierr = PetscCalloc1(c->m,&rowhit);CHKERRQ(ierr); 233 ierr = PetscMalloc1(c->m,&valaddrhit);CHKERRQ(ierr); 234 235 nz = 0; 236 for (i=0; i<nis; i++) { /* loop over colors */ 237 ierr = ISGetLocalSize(c->isa[i],&n);CHKERRQ(ierr); 238 ierr = ISGetIndices(c->isa[i],&is);CHKERRQ(ierr); 239 240 c->ncolumns[i] = n; 241 c->columns[i] = (PetscInt*)is; 242 /* note: we know that c->isa is going to be around as long at the c->columns values */ 243 ierr = ISRestoreIndices(c->isa[i],&is);CHKERRQ(ierr); 244 245 /* fast, crude version requires O(N*N) work */ 246 bs2 = bs*bs; 247 nrows = 0; 248 for (j=0; j<n; j++) { /* loop over columns */ 249 col = is[j]; 250 tmp = ci[col]; 251 row = cj + tmp; 252 m = ci[col+1] - tmp; 253 nrows += m; 254 for (k=0; k<m; k++) { /* loop over columns marking them in rowhit */ 255 rowhit[*row] = col + 1; 256 valaddrhit[*row++] = (PetscScalar*)&A_val[bs2*spidx[tmp + k]]; 257 } 258 } 259 c->nrows[i] = nrows; /* total num of rows for this color */ 260 261 if (c->htype[0] == 'd') { 262 for (j=0; j<mbs; j++) { /* loop over rows */ 263 if (rowhit[j]) { 264 Jentry[nz].row = j; /* local row index */ 265 Jentry[nz].col = rowhit[j] - 1; /* local column index */ 266 Jentry[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */ 267 nz++; 268 rowhit[j] = 0.0; /* zero rowhit for reuse */ 269 } 270 } 271 } else { /* c->htype == 'wp' */ 272 for (j=0; j<mbs; j++) { /* loop over rows */ 273 if (rowhit[j]) { 274 Jentry2[nz].row = j; /* local row index */ 275 Jentry2[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */ 276 nz++; 277 rowhit[j] = 0.0; /* zero rowhit for reuse */ 278 } 279 } 280 } 281 } 282 283 if (c->bcols > 1) { /* reorder Jentry for faster MatFDColoringApply() */ 284 ierr = MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz);CHKERRQ(ierr); 285 } 286 287 if (isBAIJ) { 288 ierr = MatRestoreColumnIJ_SeqBAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr); 289 ierr = PetscMalloc1(bs*mat->rmap->n,&c->dy);CHKERRQ(ierr); 290 ierr = PetscLogObjectMemory((PetscObject)c,bs*mat->rmap->n*sizeof(PetscScalar));CHKERRQ(ierr); 291 } else if (isSELL) { 292 ierr = MatRestoreColumnIJ_SeqSELL_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr); 293 } else { 294 ierr = MatRestoreColumnIJ_SeqAIJ_Color(mat,0,PETSC_FALSE,PETSC_FALSE,&ncols,&ci,&cj,&spidx,NULL);CHKERRQ(ierr); 295 } 296 ierr = PetscFree(rowhit);CHKERRQ(ierr); 297 ierr = PetscFree(valaddrhit);CHKERRQ(ierr); 298 ierr = ISColoringRestoreIS(iscoloring,PETSC_OWN_POINTER,&c->isa);CHKERRQ(ierr); 299 300 ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->rmap->n,PETSC_DETERMINE,0,NULL,&c->vscale);CHKERRQ(ierr); 301 ierr = PetscInfo3(c,"ncolors %D, brows %D and bcols %D are used.\n",c->ncolors,c->brows,c->bcols);CHKERRQ(ierr); 302 PetscFunctionReturn(0); 303 } 304