1 2 #include <../src/mat/impls/aij/mpi/mpiaij.h> 3 #include <../src/mat/impls/baij/mpi/mpibaij.h> 4 5 #undef __FUNCT__ 6 #define __FUNCT__ "MatFDColoringApply_BAIJ" 7 PetscErrorCode MatFDColoringApply_BAIJ(Mat J,MatFDColoring coloring,Vec x1,MatStructure *flag,void *sctx) 8 { 9 PetscErrorCode (*f)(void*,Vec,Vec,void*) = (PetscErrorCode (*)(void*,Vec,Vec,void*))coloring->f; 10 PetscErrorCode ierr; 11 PetscInt k,cstart,cend,l,row,col,nz,spidx,i,j; 12 PetscScalar dx=0.0,*y,*xx,*w3_array,*dy_i,*dy=coloring->dy; 13 PetscScalar *vscale_array; 14 PetscReal epsilon=coloring->error_rel,umin=coloring->umin,unorm; 15 Vec w1=coloring->w1,w2=coloring->w2,w3,vscale=coloring->vscale; 16 void *fctx=coloring->fctx; 17 PetscBool flg=PETSC_FALSE; 18 PetscInt ctype=coloring->ctype,nxloc,nrows_k; 19 PetscScalar *valaddr; 20 MatEntry *Jentry=coloring->matentry; 21 const PetscInt ncolors=coloring->ncolors,*ncolumns=coloring->ncolumns,*nrows=coloring->nrows; 22 PetscInt bs=J->rmap->bs; 23 24 PetscFunctionBegin; 25 ierr = MatSetUnfactored(J);CHKERRQ(ierr); 26 ierr = PetscOptionsGetBool(NULL,"-mat_fd_coloring_dont_rezero",&flg,NULL);CHKERRQ(ierr); 27 if (flg) { 28 ierr = PetscInfo(coloring,"Not calling MatZeroEntries()\n");CHKERRQ(ierr); 29 } else { 30 PetscBool assembled; 31 ierr = MatAssembled(J,&assembled);CHKERRQ(ierr); 32 if (assembled) { 33 ierr = MatZeroEntries(J);CHKERRQ(ierr); 34 } 35 } 36 37 /* create vscale for storing dx */ 38 if (!vscale) { 39 if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') { 40 /* contain the "diagonal" on processor scalings followed by the off processor* - garray must be non-bloced! */ 41 Mat_MPIBAIJ *baij=(Mat_MPIBAIJ*)J->data; 42 PetscInt *garray; 43 ierr = PetscMalloc(baij->B->cmap->n*sizeof(PetscInt),&garray);CHKERRQ(ierr); 44 for (i=0; i<baij->B->cmap->n/bs; i++) { 45 for (j=0; j<bs; j++) { 46 garray[i*bs+j] = bs*baij->garray[i]+j; 47 } 48 } 49 ierr = VecCreateGhost(PetscObjectComm((PetscObject)J),J->cmap->n,PETSC_DETERMINE,baij->B->cmap->n,garray,&vscale);CHKERRQ(ierr); 50 ierr = PetscFree(garray);CHKERRQ(ierr); 51 } else if (ctype == IS_COLORING_GHOSTED) { 52 ierr = VecDuplicate(x1,&vscale);CHKERRQ(ierr); 53 } 54 coloring->vscale = vscale; 55 } 56 57 /* (1) Set w1 = F(x1) */ 58 if (!coloring->fset) { 59 ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 60 ierr = (*f)(sctx,x1,w1,fctx);CHKERRQ(ierr); 61 ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 62 } else { 63 coloring->fset = PETSC_FALSE; 64 } 65 66 /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */ 67 ierr = VecGetLocalSize(x1,&nxloc);CHKERRQ(ierr); 68 if (coloring->htype[0] == 'w') { 69 /* vscale = dx is a constant scalar */ 70 ierr = VecNorm(x1,NORM_2,&unorm);CHKERRQ(ierr); 71 dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon); 72 } else { 73 ierr = VecGetArray(x1,&xx);CHKERRQ(ierr); 74 ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr); 75 for (col=0; col<nxloc; col++) { 76 dx = xx[col]; 77 if (PetscAbsScalar(dx) < umin) { 78 if (PetscRealPart(dx) >= 0.0) dx = umin; 79 else if (PetscRealPart(dx) < 0.0 ) dx = -umin; 80 } 81 dx *= epsilon; 82 vscale_array[col] = 1.0/dx; 83 } 84 ierr = VecRestoreArray(x1,&xx);CHKERRQ(ierr); 85 ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr); 86 } 87 if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] != 'w') { 88 ierr = VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr); 89 ierr = VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr); 90 } 91 92 /* (3) Loop over each color */ 93 if (!coloring->w3) { 94 ierr = VecDuplicate(x1,&coloring->w3);CHKERRQ(ierr); 95 ierr = PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);CHKERRQ(ierr); 96 } 97 w3 = coloring->w3; 98 99 ierr = VecGetOwnershipRange(x1,&cstart,&cend);CHKERRQ(ierr); /* used by ghosted vscale */ 100 if (vscale) { 101 ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr); 102 } 103 nz = 0; 104 for (k=0; k<ncolors; k++) { 105 coloring->currentcolor = k; 106 107 /* 108 (3-1) Loop over each column associated with color 109 adding the perturbation to the vector w3 = x1 + dx. 110 */ 111 ierr = VecCopy(x1,w3);CHKERRQ(ierr); 112 dy_i = dy; 113 for (i=0; i<bs; i++) { /* Loop over a block of columns */ 114 ierr = VecGetArray(w3,&w3_array);CHKERRQ(ierr); 115 if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */ 116 if (coloring->htype[0] == 'w') { 117 for (l=0; l<ncolumns[k]; l++) { 118 col = i + bs*coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */ 119 w3_array[col] += 1.0/dx; 120 121 if (i) { 122 w3_array[col-1] -= 1.0/dx; /* resume original w3[col-1] */ 123 } 124 125 } 126 } else { /* htype == 'ds' */ 127 vscale_array -= cstart; /* shift pointer so global index can be used */ 128 for (l=0; l<ncolumns[k]; l++) { 129 col = i + bs*coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */ 130 w3_array[col] += 1.0/vscale_array[col]; 131 132 if (i) { 133 w3_array[col-1] -= 1.0/vscale_array[col-1]; /* resume original w3[col-1] */ 134 } 135 136 } 137 vscale_array += cstart; 138 } 139 if (ctype == IS_COLORING_GLOBAL) w3_array += cstart; 140 ierr = VecRestoreArray(w3,&w3_array);CHKERRQ(ierr); 141 142 /* 143 (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations) 144 w2 = F(x1 + dx) - F(x1) 145 */ 146 ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 147 ierr = VecPlaceArray(w2,dy_i);CHKERRQ(ierr); /* place w2 to the array dy_i */ 148 ierr = (*f)(sctx,w3,w2,fctx);CHKERRQ(ierr); 149 ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 150 ierr = VecAXPY(w2,-1.0,w1);CHKERRQ(ierr); 151 ierr = VecResetArray(w2);CHKERRQ(ierr); 152 dy_i += nxloc; /* points to dy+i*nxloc */ 153 } 154 155 /* 156 (3-3) Loop over rows of vector, putting results into Jacobian matrix 157 */ 158 nrows_k = nrows[k]; 159 ierr = VecGetArray(w2,&y);CHKERRQ(ierr); 160 if (coloring->htype[0] == 'w') { 161 for (l=0; l<nrows_k; l++) { 162 row = bs*Jentry[nz].row; /* local row index */ 163 valaddr = Jentry[nz].valaddr; 164 nz++; 165 spidx = 0; 166 dy_i = dy; 167 for (i=0; i<bs; i++) { /* column of the block */ 168 for (j=0; j<bs; j++) { /* row of the block */ 169 valaddr[spidx++] = dy_i[row+j]*dx; 170 } 171 dy_i += nxloc; /* points to dy+i*nxloc */ 172 } 173 } 174 } else { /* htype == 'ds' */ 175 for (l=0; l<nrows_k; l++) { 176 row = bs*Jentry[nz].row; /* local row index */ 177 col = bs*Jentry[nz].col; /* local column index */ 178 valaddr = Jentry[nz].valaddr; 179 nz++; 180 spidx = 0; 181 dy_i = dy; 182 for (i=0; i<bs; i++) { /* column of the block */ 183 for (j=0; j<bs; j++) { /* row of the block */ 184 valaddr[spidx++] = dy_i[row+j]*vscale_array[col+i]; 185 } 186 dy_i += nxloc; /* points to dy+i*nxloc */ 187 } 188 } 189 } 190 ierr = VecRestoreArray(w2,&y);CHKERRQ(ierr); 191 } 192 ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 193 ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 194 if (vscale) { 195 ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr); 196 } 197 198 coloring->currentcolor = -1; 199 PetscFunctionReturn(0); 200 } 201 202 #undef __FUNCT__ 203 #define __FUNCT__ "MatFDColoringApply_AIJ" 204 PetscErrorCode MatFDColoringApply_AIJ(Mat J,MatFDColoring coloring,Vec x1,MatStructure *flag,void *sctx) 205 { 206 PetscErrorCode (*f)(void*,Vec,Vec,void*) = (PetscErrorCode (*)(void*,Vec,Vec,void*))coloring->f; 207 PetscErrorCode ierr; 208 PetscInt k,cstart,cend,l,row,col,nz; 209 PetscScalar dx=0.0,*y,*xx,*w3_array; 210 PetscScalar *vscale_array; 211 PetscReal epsilon=coloring->error_rel,umin=coloring->umin,unorm; 212 Vec w1=coloring->w1,w2=coloring->w2,w3,vscale=coloring->vscale; 213 void *fctx=coloring->fctx; 214 PetscBool flg=PETSC_FALSE; 215 PetscInt ctype=coloring->ctype,nxloc,nrows_k; 216 MatEntry *Jentry=coloring->matentry; 217 const PetscInt ncolors=coloring->ncolors,*ncolumns=coloring->ncolumns,*nrows=coloring->nrows; 218 219 PetscFunctionBegin; 220 ierr = MatSetUnfactored(J);CHKERRQ(ierr); 221 ierr = PetscOptionsGetBool(NULL,"-mat_fd_coloring_dont_rezero",&flg,NULL);CHKERRQ(ierr); 222 if (flg) { 223 ierr = PetscInfo(coloring,"Not calling MatZeroEntries()\n");CHKERRQ(ierr); 224 } else { 225 PetscBool assembled; 226 ierr = MatAssembled(J,&assembled);CHKERRQ(ierr); 227 if (assembled) { 228 ierr = MatZeroEntries(J);CHKERRQ(ierr); 229 } 230 } 231 232 /* create vscale for storing dx */ 233 if (!vscale) { 234 if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') { 235 Mat_MPIAIJ *aij=(Mat_MPIAIJ*)J->data; 236 ierr = VecCreateGhost(PetscObjectComm((PetscObject)J),J->cmap->n,PETSC_DETERMINE,aij->B->cmap->n,aij->garray,&vscale);CHKERRQ(ierr); 237 } else if (ctype == IS_COLORING_GHOSTED) { 238 ierr = VecDuplicate(x1,&vscale);CHKERRQ(ierr); 239 } 240 coloring->vscale = vscale; 241 } 242 243 /* (1) Set w1 = F(x1) */ 244 if (!coloring->fset) { 245 ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 246 ierr = (*f)(sctx,x1,w1,fctx);CHKERRQ(ierr); 247 ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 248 } else { 249 coloring->fset = PETSC_FALSE; 250 } 251 252 /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */ 253 if (coloring->htype[0] == 'w') { 254 /* vscale = dx is a constant scalar */ 255 ierr = VecNorm(x1,NORM_2,&unorm);CHKERRQ(ierr); 256 dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon); 257 } else { 258 ierr = VecGetLocalSize(x1,&nxloc);CHKERRQ(ierr); 259 ierr = VecGetArray(x1,&xx);CHKERRQ(ierr); 260 ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr); 261 for (col=0; col<nxloc; col++) { 262 dx = xx[col]; 263 if (PetscAbsScalar(dx) < umin) { 264 if (PetscRealPart(dx) >= 0.0) dx = umin; 265 else if (PetscRealPart(dx) < 0.0 ) dx = -umin; 266 } 267 dx *= epsilon; 268 vscale_array[col] = 1.0/dx; 269 } 270 ierr = VecRestoreArray(x1,&xx);CHKERRQ(ierr); 271 ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr); 272 } 273 if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] != 'w') { 274 ierr = VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr); 275 ierr = VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr); 276 } 277 278 /* (3) Loop over each color */ 279 if (!coloring->w3) { 280 ierr = VecDuplicate(x1,&coloring->w3);CHKERRQ(ierr); 281 ierr = PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);CHKERRQ(ierr); 282 } 283 w3 = coloring->w3; 284 285 ierr = VecGetOwnershipRange(x1,&cstart,&cend);CHKERRQ(ierr); /* used by ghosted vscale */ 286 if (vscale) { 287 ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr); 288 } 289 nz = 0; 290 for (k=0; k<ncolors; k++) { 291 coloring->currentcolor = k; 292 293 /* 294 (3-1) Loop over each column associated with color 295 adding the perturbation to the vector w3 = x1 + dx. 296 */ 297 ierr = VecCopy(x1,w3);CHKERRQ(ierr); 298 ierr = VecGetArray(w3,&w3_array);CHKERRQ(ierr); 299 if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */ 300 if (coloring->htype[0] == 'w') { 301 for (l=0; l<ncolumns[k]; l++) { 302 col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */ 303 w3_array[col] += 1.0/dx; 304 } 305 } else { /* htype == 'ds' */ 306 vscale_array -= cstart; /* shift pointer so global index can be used */ 307 for (l=0; l<ncolumns[k]; l++) { 308 col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */ 309 w3_array[col] += 1.0/vscale_array[col]; 310 } 311 vscale_array += cstart; 312 } 313 if (ctype == IS_COLORING_GLOBAL) w3_array += cstart; 314 ierr = VecRestoreArray(w3,&w3_array);CHKERRQ(ierr); 315 316 /* 317 (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations) 318 w2 = F(x1 + dx) - F(x1) 319 */ 320 ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 321 ierr = (*f)(sctx,w3,w2,fctx);CHKERRQ(ierr); 322 ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 323 ierr = VecAXPY(w2,-1.0,w1);CHKERRQ(ierr); 324 325 /* 326 (3-3) Loop over rows of vector, putting results into Jacobian matrix 327 */ 328 nrows_k = nrows[k]; 329 ierr = VecGetArray(w2,&y);CHKERRQ(ierr); 330 if (coloring->htype[0] == 'w') { 331 for (l=0; l<nrows_k; l++) { 332 row = Jentry[nz].row; /* local row index */ 333 *(Jentry[nz++].valaddr) = y[row]*dx; 334 } 335 } else { /* htype == 'ds' */ 336 for (l=0; l<nrows_k; l++) { 337 row = Jentry[nz].row; /* local row index */ 338 *(Jentry[nz].valaddr) = y[row]*vscale_array[Jentry[nz].col]; 339 nz++; 340 } 341 } 342 ierr = VecRestoreArray(w2,&y);CHKERRQ(ierr); 343 } 344 ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 345 ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 346 if (vscale) { 347 ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr); 348 } 349 350 coloring->currentcolor = -1; 351 PetscFunctionReturn(0); 352 } 353 354 #undef __FUNCT__ 355 #define __FUNCT__ "MatFDColoringCreate_MPIXAIJ" 356 PetscErrorCode MatFDColoringCreate_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c) 357 { 358 PetscErrorCode ierr; 359 PetscMPIInt size,*ncolsonproc,*disp,nn; 360 PetscInt i,n,nrows,nrows_i,j,k,m,ncols,col,*rowhit,cstart,cend,colb; 361 const PetscInt *is,*A_ci,*A_cj,*B_ci,*B_cj,*row=NULL,*ltog=NULL; 362 PetscInt nis=iscoloring->n,nctot,*cols; 363 IS *isa; 364 ISLocalToGlobalMapping map=mat->cmap->mapping; 365 PetscInt ctype=c->ctype,*spidxA,*spidxB,nz,bs,bs2,spidx; 366 Mat A,B; 367 PetscScalar *A_val,*B_val,**valaddrhit; 368 MatEntry *Jentry; 369 PetscBool isBAIJ; 370 #if defined(PETSC_USE_CTABLE) 371 PetscTable colmap=NULL; 372 #else 373 PetscInt *colmap=NULL; /* local col number of off-diag col */ 374 #endif 375 376 PetscFunctionBegin; 377 if (ctype == IS_COLORING_GHOSTED) { 378 if (!map) SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_INCOMP,"When using ghosted differencing matrix must have local to global mapping provided with MatSetLocalToGlobalMapping"); 379 ierr = ISLocalToGlobalMappingGetIndices(map,<og);CHKERRQ(ierr); 380 } 381 382 ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr); 383 ierr = PetscObjectTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);CHKERRQ(ierr); 384 if (isBAIJ) { 385 Mat_MPIBAIJ *aij=(Mat_MPIBAIJ*)mat->data; 386 Mat_SeqBAIJ *spA,*spB; 387 A = aij->A; spA = (Mat_SeqBAIJ*)A->data; A_val = spA->a; 388 B = aij->B; spB = (Mat_SeqBAIJ*)B->data; B_val = spB->a; 389 nz = spA->nz + spB->nz; /* total nonzero entries of mat */ 390 if (!aij->colmap) { 391 ierr = MatCreateColmap_MPIBAIJ_Private(mat);CHKERRQ(ierr); 392 colmap = aij->colmap; 393 } 394 ierr = MatGetColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr); 395 ierr = MatGetColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr); 396 } else { 397 Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data; 398 Mat_SeqAIJ *spA,*spB; 399 A = aij->A; spA = (Mat_SeqAIJ*)A->data; A_val = spA->a; 400 B = aij->B; spB = (Mat_SeqAIJ*)B->data; B_val = spB->a; 401 nz = spA->nz + spB->nz; /* total nonzero entries of mat */ 402 if (!aij->colmap) { 403 /* Allow access to data structures of local part of matrix 404 - creates aij->colmap which maps global column number to local number in part B */ 405 ierr = MatCreateColmap_MPIAIJ_Private(mat);CHKERRQ(ierr); 406 colmap = aij->colmap; 407 } 408 ierr = MatGetColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr); 409 ierr = MatGetColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr); 410 411 bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */ 412 } 413 414 m = mat->rmap->n/bs; 415 cstart = mat->cmap->rstart/bs; 416 cend = mat->cmap->rend/bs; 417 c->M = mat->rmap->N/bs; /* set the global rows and columns and local rows */ 418 c->N = mat->cmap->N/bs; 419 c->m = m; 420 c->rstart = mat->rmap->rstart/bs; 421 422 c->ncolors = nis; 423 ierr = PetscMalloc(nis*sizeof(PetscInt),&c->ncolumns);CHKERRQ(ierr); 424 ierr = PetscMalloc(nis*sizeof(PetscInt*),&c->columns);CHKERRQ(ierr); 425 ierr = PetscMalloc(nis*sizeof(PetscInt),&c->nrows);CHKERRQ(ierr); 426 ierr = PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));CHKERRQ(ierr); 427 428 ierr = PetscMalloc(nz*sizeof(MatEntry),&Jentry);CHKERRQ(ierr); 429 ierr = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));CHKERRQ(ierr); 430 c->matentry = Jentry; 431 432 ierr = PetscMalloc2(m+1,PetscInt,&rowhit,m+1,PetscScalar*,&valaddrhit);CHKERRQ(ierr); 433 nz = 0; 434 ierr = ISColoringGetIS(iscoloring,PETSC_IGNORE,&isa);CHKERRQ(ierr); 435 for (i=0; i<nis; i++) { /* for each local color */ 436 ierr = ISGetLocalSize(isa[i],&n);CHKERRQ(ierr); 437 ierr = ISGetIndices(isa[i],&is);CHKERRQ(ierr); 438 439 c->ncolumns[i] = n; /* local number of columns of this color on this process */ 440 if (n) { 441 ierr = PetscMalloc(n*sizeof(PetscInt),&c->columns[i]);CHKERRQ(ierr); 442 ierr = PetscLogObjectMemory((PetscObject)c,n*sizeof(PetscInt));CHKERRQ(ierr); 443 ierr = PetscMemcpy(c->columns[i],is,n*sizeof(PetscInt));CHKERRQ(ierr); 444 } else { 445 c->columns[i] = 0; 446 } 447 448 if (ctype == IS_COLORING_GLOBAL) { 449 /* Determine nctot, the total (parallel) number of columns of this color */ 450 ierr = MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size);CHKERRQ(ierr); 451 ierr = PetscMalloc2(size,PetscMPIInt,&ncolsonproc,size,PetscMPIInt,&disp);CHKERRQ(ierr); 452 453 /* ncolsonproc[j]: local ncolumns on proc[j] of this color */ 454 ierr = PetscMPIIntCast(n,&nn);CHKERRQ(ierr); 455 ierr = MPI_Allgather(&nn,1,MPI_INT,ncolsonproc,1,MPI_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr); 456 nctot = 0; for (j=0; j<size; j++) nctot += ncolsonproc[j]; 457 if (!nctot) { 458 ierr = PetscInfo(mat,"Coloring of matrix has some unneeded colors with no corresponding rows\n");CHKERRQ(ierr); 459 } 460 461 disp[0] = 0; 462 for (j=1; j<size; j++) { 463 disp[j] = disp[j-1] + ncolsonproc[j-1]; 464 } 465 466 /* Get cols, the complete list of columns for this color on each process */ 467 ierr = PetscMalloc((nctot+1)*sizeof(PetscInt),&cols);CHKERRQ(ierr); 468 ierr = MPI_Allgatherv((void*)is,n,MPIU_INT,cols,ncolsonproc,disp,MPIU_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr); 469 ierr = PetscFree2(ncolsonproc,disp);CHKERRQ(ierr); 470 } else if (ctype == IS_COLORING_GHOSTED) { 471 /* Determine local number of columns of this color on this process, including ghost points */ 472 nctot = n; 473 ierr = PetscMalloc((nctot+1)*sizeof(PetscInt),&cols);CHKERRQ(ierr); 474 ierr = PetscMemcpy(cols,is,n*sizeof(PetscInt));CHKERRQ(ierr); 475 } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not provided for this MatFDColoring type"); 476 477 /* Mark all rows affect by these columns */ 478 ierr = PetscMemzero(rowhit,m*sizeof(PetscInt));CHKERRQ(ierr); 479 bs2 = bs*bs; 480 nrows_i = 0; 481 for (j=0; j<nctot; j++) { /* loop over columns*/ 482 if (ctype == IS_COLORING_GHOSTED) { 483 col = ltog[cols[j]]; 484 } else { 485 col = cols[j]; 486 } 487 if (col >= cstart && col < cend) { /* column is in A, diagonal block of mat */ 488 row = A_cj + A_ci[col-cstart]; 489 nrows = A_ci[col-cstart+1] - A_ci[col-cstart]; 490 nrows_i += nrows; 491 /* loop over columns of A marking them in rowhit */ 492 for (k=0; k<nrows; k++) { 493 /* set valaddrhit for part A */ 494 spidx = bs2*spidxA[A_ci[col-cstart] + k]; 495 valaddrhit[*row] = &A_val[spidx]; 496 rowhit[*row++] = col - cstart + 1; /* local column index */ 497 } 498 } else { /* column is in B, off-diagonal block of mat */ 499 #if defined(PETSC_USE_CTABLE) 500 ierr = PetscTableFind(colmap,col+1,&colb);CHKERRQ(ierr); 501 colb--; 502 #else 503 colb = colmap[col] - 1; /* local column index */ 504 #endif 505 if (colb == -1) { 506 nrows = 0; 507 } else { 508 colb = colb/bs; 509 row = B_cj + B_ci[colb]; 510 nrows = B_ci[colb+1] - B_ci[colb]; 511 } 512 nrows_i += nrows; 513 /* loop over columns of B marking them in rowhit */ 514 for (k=0; k<nrows; k++) { 515 /* set valaddrhit for part B */ 516 spidx = bs2*spidxB[B_ci[colb] + k]; 517 valaddrhit[*row] = &B_val[spidx]; 518 rowhit[*row++] = colb + 1 + cend - cstart; /* local column index */ 519 } 520 } 521 } 522 c->nrows[i] = nrows_i; 523 524 for (j=0; j<m; j++) { 525 if (rowhit[j]) { 526 Jentry[nz].row = j; /* local row index */ 527 Jentry[nz].col = rowhit[j] - 1; /* local column index */ 528 Jentry[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */ 529 nz++; 530 } 531 } 532 ierr = PetscFree(cols);CHKERRQ(ierr); 533 } 534 ierr = ISColoringRestoreIS(iscoloring,&isa);CHKERRQ(ierr); 535 536 ierr = PetscFree2(rowhit,valaddrhit);CHKERRQ(ierr); 537 if (isBAIJ) { 538 ierr = MatRestoreColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr); 539 ierr = MatRestoreColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr); 540 ierr = PetscMalloc(bs*mat->rmap->n*sizeof(PetscScalar),&c->dy);CHKERRQ(ierr); 541 } else { 542 ierr = MatRestoreColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr); 543 ierr = MatRestoreColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr); 544 } 545 546 if (ctype == IS_COLORING_GHOSTED) { 547 ierr = ISLocalToGlobalMappingRestoreIndices(map,<og);CHKERRQ(ierr); 548 } 549 PetscFunctionReturn(0); 550 } 551