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