1 #include <../src/mat/impls/sell/mpi/mpisell.h> 2 #include <../src/mat/impls/aij/mpi/mpiaij.h> 3 #include <../src/mat/impls/baij/mpi/mpibaij.h> 4 #include <petsc/private/isimpl.h> 5 6 PetscErrorCode MatFDColoringApply_BAIJ(Mat J,MatFDColoring coloring,Vec x1,void *sctx) 7 { 8 PetscErrorCode (*f)(void*,Vec,Vec,void*)=(PetscErrorCode (*)(void*,Vec,Vec,void*))coloring->f; 9 PetscErrorCode ierr; 10 PetscInt k,cstart,cend,l,row,col,nz,spidx,i,j; 11 PetscScalar dx=0.0,*w3_array,*dy_i,*dy=coloring->dy; 12 PetscScalar *vscale_array; 13 const PetscScalar *xx; 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 PetscInt ctype=coloring->ctype,nxloc,nrows_k; 18 PetscScalar *valaddr; 19 MatEntry *Jentry=coloring->matentry; 20 MatEntry2 *Jentry2=coloring->matentry2; 21 const PetscInt ncolors=coloring->ncolors,*ncolumns=coloring->ncolumns,*nrows=coloring->nrows; 22 PetscInt bs=J->rmap->bs; 23 24 PetscFunctionBegin; 25 ierr = VecPinToCPU(x1,PETSC_TRUE);CHKERRQ(ierr); 26 /* (1) Set w1 = F(x1) */ 27 if (!coloring->fset) { 28 ierr = PetscLogEventBegin(MAT_FDColoringFunction,coloring,0,0,0);CHKERRQ(ierr); 29 ierr = (*f)(sctx,x1,w1,fctx);CHKERRQ(ierr); 30 ierr = PetscLogEventEnd(MAT_FDColoringFunction,coloring,0,0,0);CHKERRQ(ierr); 31 } else { 32 coloring->fset = PETSC_FALSE; 33 } 34 35 /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */ 36 ierr = VecGetLocalSize(x1,&nxloc);CHKERRQ(ierr); 37 if (coloring->htype[0] == 'w') { 38 /* vscale = dx is a constant scalar */ 39 ierr = VecNorm(x1,NORM_2,&unorm);CHKERRQ(ierr); 40 dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon); 41 } else { 42 ierr = VecGetArrayRead(x1,&xx);CHKERRQ(ierr); 43 ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr); 44 for (col=0; col<nxloc; col++) { 45 dx = xx[col]; 46 if (PetscAbsScalar(dx) < umin) { 47 if (PetscRealPart(dx) >= 0.0) dx = umin; 48 else if (PetscRealPart(dx) < 0.0 ) dx = -umin; 49 } 50 dx *= epsilon; 51 vscale_array[col] = 1.0/dx; 52 } 53 ierr = VecRestoreArrayRead(x1,&xx);CHKERRQ(ierr); 54 ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr); 55 } 56 if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') { 57 ierr = VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr); 58 ierr = VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr); 59 } 60 61 /* (3) Loop over each color */ 62 if (!coloring->w3) { 63 ierr = VecDuplicate(x1,&coloring->w3);CHKERRQ(ierr); 64 /* Vec is used instensively in particular piece of scalar CPU code; won't benifit from bouncing back and forth to the GPU */ 65 ierr = VecPinToCPU(coloring->w3,PETSC_TRUE);CHKERRQ(ierr); 66 ierr = PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);CHKERRQ(ierr); 67 } 68 w3 = coloring->w3; 69 70 ierr = VecGetOwnershipRange(x1,&cstart,&cend);CHKERRQ(ierr); /* used by ghosted vscale */ 71 if (vscale) { 72 ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr); 73 } 74 nz = 0; 75 for (k=0; k<ncolors; k++) { 76 coloring->currentcolor = k; 77 78 /* 79 (3-1) Loop over each column associated with color 80 adding the perturbation to the vector w3 = x1 + dx. 81 */ 82 ierr = VecCopy(x1,w3);CHKERRQ(ierr); 83 dy_i = dy; 84 for (i=0; i<bs; i++) { /* Loop over a block of columns */ 85 ierr = VecGetArray(w3,&w3_array);CHKERRQ(ierr); 86 if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */ 87 if (coloring->htype[0] == 'w') { 88 for (l=0; l<ncolumns[k]; l++) { 89 col = i + bs*coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */ 90 w3_array[col] += 1.0/dx; 91 if (i) w3_array[col-1] -= 1.0/dx; /* resume original w3[col-1] */ 92 } 93 } else { /* htype == 'ds' */ 94 vscale_array -= cstart; /* shift pointer so global index can be used */ 95 for (l=0; l<ncolumns[k]; l++) { 96 col = i + bs*coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */ 97 w3_array[col] += 1.0/vscale_array[col]; 98 if (i) w3_array[col-1] -= 1.0/vscale_array[col-1]; /* resume original w3[col-1] */ 99 } 100 vscale_array += cstart; 101 } 102 if (ctype == IS_COLORING_GLOBAL) w3_array += cstart; 103 ierr = VecRestoreArray(w3,&w3_array);CHKERRQ(ierr); 104 105 /* 106 (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations) 107 w2 = F(x1 + dx) - F(x1) 108 */ 109 ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 110 ierr = VecPlaceArray(w2,dy_i);CHKERRQ(ierr); /* place w2 to the array dy_i */ 111 ierr = (*f)(sctx,w3,w2,fctx);CHKERRQ(ierr); 112 ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 113 ierr = VecAXPY(w2,-1.0,w1);CHKERRQ(ierr); 114 ierr = VecResetArray(w2);CHKERRQ(ierr); 115 dy_i += nxloc; /* points to dy+i*nxloc */ 116 } 117 118 /* 119 (3-3) Loop over rows of vector, putting results into Jacobian matrix 120 */ 121 nrows_k = nrows[k]; 122 if (coloring->htype[0] == 'w') { 123 for (l=0; l<nrows_k; l++) { 124 row = bs*Jentry2[nz].row; /* local row index */ 125 valaddr = Jentry2[nz++].valaddr; 126 spidx = 0; 127 dy_i = dy; 128 for (i=0; i<bs; i++) { /* column of the block */ 129 for (j=0; j<bs; j++) { /* row of the block */ 130 valaddr[spidx++] = dy_i[row+j]*dx; 131 } 132 dy_i += nxloc; /* points to dy+i*nxloc */ 133 } 134 } 135 } else { /* htype == 'ds' */ 136 for (l=0; l<nrows_k; l++) { 137 row = bs*Jentry[nz].row; /* local row index */ 138 col = bs*Jentry[nz].col; /* local column index */ 139 valaddr = Jentry[nz++].valaddr; 140 spidx = 0; 141 dy_i = dy; 142 for (i=0; i<bs; i++) { /* column of the block */ 143 for (j=0; j<bs; j++) { /* row of the block */ 144 valaddr[spidx++] = dy_i[row+j]*vscale_array[col+i]; 145 } 146 dy_i += nxloc; /* points to dy+i*nxloc */ 147 } 148 } 149 } 150 } 151 ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 152 ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 153 if (vscale) { 154 ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr); 155 } 156 157 coloring->currentcolor = -1; 158 ierr = VecPinToCPU(x1,PETSC_FALSE);CHKERRQ(ierr); 159 PetscFunctionReturn(0); 160 } 161 162 /* this is declared PETSC_EXTERN because it is used by MatFDColoringUseDM() which is in the DM library */ 163 PetscErrorCode MatFDColoringApply_AIJ(Mat J,MatFDColoring coloring,Vec x1,void *sctx) 164 { 165 PetscErrorCode (*f)(void*,Vec,Vec,void*) = (PetscErrorCode (*)(void*,Vec,Vec,void*))coloring->f; 166 PetscErrorCode ierr; 167 PetscInt k,cstart,cend,l,row,col,nz; 168 PetscScalar dx=0.0,*y,*w3_array; 169 const PetscScalar *xx; 170 PetscScalar *vscale_array; 171 PetscReal epsilon=coloring->error_rel,umin=coloring->umin,unorm; 172 Vec w1=coloring->w1,w2=coloring->w2,w3,vscale=coloring->vscale; 173 void *fctx=coloring->fctx; 174 ISColoringType ctype=coloring->ctype; 175 PetscInt nxloc,nrows_k; 176 MatEntry *Jentry=coloring->matentry; 177 MatEntry2 *Jentry2=coloring->matentry2; 178 const PetscInt ncolors=coloring->ncolors,*ncolumns=coloring->ncolumns,*nrows=coloring->nrows; 179 180 PetscFunctionBegin; 181 ierr = VecPinToCPU(x1,PETSC_TRUE);CHKERRQ(ierr); 182 if ((ctype == IS_COLORING_LOCAL) && (J->ops->fdcoloringapply == MatFDColoringApply_AIJ)) SETERRQ(PetscObjectComm((PetscObject)J),PETSC_ERR_SUP,"Must call MatColoringUseDM() with IS_COLORING_LOCAL"); 183 /* (1) Set w1 = F(x1) */ 184 if (!coloring->fset) { 185 ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 186 ierr = (*f)(sctx,x1,w1,fctx);CHKERRQ(ierr); 187 ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 188 } else { 189 coloring->fset = PETSC_FALSE; 190 } 191 192 /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */ 193 if (coloring->htype[0] == 'w') { 194 /* vscale = 1./dx is a constant scalar */ 195 ierr = VecNorm(x1,NORM_2,&unorm);CHKERRQ(ierr); 196 dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon); 197 } else { 198 ierr = VecGetLocalSize(x1,&nxloc);CHKERRQ(ierr); 199 ierr = VecGetArrayRead(x1,&xx);CHKERRQ(ierr); 200 ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr); 201 for (col=0; col<nxloc; col++) { 202 dx = xx[col]; 203 if (PetscAbsScalar(dx) < umin) { 204 if (PetscRealPart(dx) >= 0.0) dx = umin; 205 else if (PetscRealPart(dx) < 0.0 ) dx = -umin; 206 } 207 dx *= epsilon; 208 vscale_array[col] = 1.0/dx; 209 } 210 ierr = VecRestoreArrayRead(x1,&xx);CHKERRQ(ierr); 211 ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr); 212 } 213 if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') { 214 ierr = VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr); 215 ierr = VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr); 216 } 217 218 /* (3) Loop over each color */ 219 if (!coloring->w3) { 220 ierr = VecDuplicate(x1,&coloring->w3);CHKERRQ(ierr); 221 ierr = PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);CHKERRQ(ierr); 222 } 223 w3 = coloring->w3; 224 225 ierr = VecGetOwnershipRange(x1,&cstart,&cend);CHKERRQ(ierr); /* used by ghosted vscale */ 226 if (vscale) { 227 ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr); 228 } 229 nz = 0; 230 231 if (coloring->bcols > 1) { /* use blocked insertion of Jentry */ 232 PetscInt i,m=J->rmap->n,nbcols,bcols=coloring->bcols; 233 PetscScalar *dy=coloring->dy,*dy_k; 234 235 nbcols = 0; 236 for (k=0; k<ncolors; k+=bcols) { 237 238 /* 239 (3-1) Loop over each column associated with color 240 adding the perturbation to the vector w3 = x1 + dx. 241 */ 242 243 dy_k = dy; 244 if (k + bcols > ncolors) bcols = ncolors - k; 245 for (i=0; i<bcols; i++) { 246 coloring->currentcolor = k+i; 247 248 ierr = VecCopy(x1,w3);CHKERRQ(ierr); 249 ierr = VecGetArray(w3,&w3_array);CHKERRQ(ierr); 250 if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */ 251 if (coloring->htype[0] == 'w') { 252 for (l=0; l<ncolumns[k+i]; l++) { 253 col = coloring->columns[k+i][l]; /* local column (in global index!) of the matrix we are probing for */ 254 w3_array[col] += 1.0/dx; 255 } 256 } else { /* htype == 'ds' */ 257 vscale_array -= cstart; /* shift pointer so global index can be used */ 258 for (l=0; l<ncolumns[k+i]; l++) { 259 col = coloring->columns[k+i][l]; /* local column (in global index!) of the matrix we are probing for */ 260 w3_array[col] += 1.0/vscale_array[col]; 261 } 262 vscale_array += cstart; 263 } 264 if (ctype == IS_COLORING_GLOBAL) w3_array += cstart; 265 ierr = VecRestoreArray(w3,&w3_array);CHKERRQ(ierr); 266 267 /* 268 (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations) 269 w2 = F(x1 + dx) - F(x1) 270 */ 271 ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 272 ierr = VecPlaceArray(w2,dy_k);CHKERRQ(ierr); /* place w2 to the array dy_i */ 273 ierr = (*f)(sctx,w3,w2,fctx);CHKERRQ(ierr); 274 ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 275 ierr = VecAXPY(w2,-1.0,w1);CHKERRQ(ierr); 276 ierr = VecResetArray(w2);CHKERRQ(ierr); 277 dy_k += m; /* points to dy+i*nxloc */ 278 } 279 280 /* 281 (3-3) Loop over block rows of vector, putting results into Jacobian matrix 282 */ 283 nrows_k = nrows[nbcols++]; 284 285 if (coloring->htype[0] == 'w') { 286 for (l=0; l<nrows_k; l++) { 287 row = Jentry2[nz].row; /* local row index */ 288 *(Jentry2[nz++].valaddr) = dy[row]*dx; 289 } 290 } else { /* htype == 'ds' */ 291 for (l=0; l<nrows_k; l++) { 292 row = Jentry[nz].row; /* local row index */ 293 *(Jentry[nz].valaddr) = dy[row]*vscale_array[Jentry[nz].col]; 294 nz++; 295 } 296 } 297 } 298 } else { /* bcols == 1 */ 299 for (k=0; k<ncolors; k++) { 300 coloring->currentcolor = k; 301 302 /* 303 (3-1) Loop over each column associated with color 304 adding the perturbation to the vector w3 = x1 + dx. 305 */ 306 ierr = VecCopy(x1,w3);CHKERRQ(ierr); 307 ierr = VecGetArray(w3,&w3_array);CHKERRQ(ierr); 308 if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */ 309 if (coloring->htype[0] == 'w') { 310 for (l=0; l<ncolumns[k]; l++) { 311 col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */ 312 w3_array[col] += 1.0/dx; 313 } 314 } else { /* htype == 'ds' */ 315 vscale_array -= cstart; /* shift pointer so global index can be used */ 316 for (l=0; l<ncolumns[k]; l++) { 317 col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */ 318 w3_array[col] += 1.0/vscale_array[col]; 319 } 320 vscale_array += cstart; 321 } 322 if (ctype == IS_COLORING_GLOBAL) w3_array += cstart; 323 ierr = VecRestoreArray(w3,&w3_array);CHKERRQ(ierr); 324 325 /* 326 (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations) 327 w2 = F(x1 + dx) - F(x1) 328 */ 329 ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 330 ierr = (*f)(sctx,w3,w2,fctx);CHKERRQ(ierr); 331 ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr); 332 ierr = VecAXPY(w2,-1.0,w1);CHKERRQ(ierr); 333 334 /* 335 (3-3) Loop over rows of vector, putting results into Jacobian matrix 336 */ 337 nrows_k = nrows[k]; 338 ierr = VecGetArray(w2,&y);CHKERRQ(ierr); 339 if (coloring->htype[0] == 'w') { 340 for (l=0; l<nrows_k; l++) { 341 row = Jentry2[nz].row; /* local row index */ 342 *(Jentry2[nz++].valaddr) = y[row]*dx; 343 } 344 } else { /* htype == 'ds' */ 345 for (l=0; l<nrows_k; l++) { 346 row = Jentry[nz].row; /* local row index */ 347 *(Jentry[nz].valaddr) = y[row]*vscale_array[Jentry[nz].col]; 348 nz++; 349 } 350 } 351 ierr = VecRestoreArray(w2,&y);CHKERRQ(ierr); 352 } 353 } 354 355 ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 356 ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 357 if (vscale) { 358 ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr); 359 } 360 coloring->currentcolor = -1; 361 ierr = VecPinToCPU(x1,PETSC_FALSE);CHKERRQ(ierr); 362 PetscFunctionReturn(0); 363 } 364 365 PetscErrorCode MatFDColoringSetUp_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c) 366 { 367 PetscErrorCode ierr; 368 PetscMPIInt size,*ncolsonproc,*disp,nn; 369 PetscInt i,n,nrows,nrows_i,j,k,m,ncols,col,*rowhit,cstart,cend,colb; 370 const PetscInt *is,*A_ci,*A_cj,*B_ci,*B_cj,*row=NULL,*ltog=NULL; 371 PetscInt nis=iscoloring->n,nctot,*cols,tmp = 0; 372 ISLocalToGlobalMapping map=mat->cmap->mapping; 373 PetscInt ctype=c->ctype,*spidxA,*spidxB,nz,bs,bs2,spidx; 374 Mat A,B; 375 PetscScalar *A_val,*B_val,**valaddrhit; 376 MatEntry *Jentry; 377 MatEntry2 *Jentry2; 378 PetscBool isBAIJ,isSELL; 379 PetscInt bcols=c->bcols; 380 #if defined(PETSC_USE_CTABLE) 381 PetscTable colmap=NULL; 382 #else 383 PetscInt *colmap=NULL; /* local col number of off-diag col */ 384 #endif 385 386 PetscFunctionBegin; 387 if (ctype == IS_COLORING_LOCAL) { 388 if (!map) SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_INCOMP,"When using ghosted differencing matrix must have local to global mapping provided with MatSetLocalToGlobalMapping"); 389 ierr = ISLocalToGlobalMappingGetIndices(map,<og);CHKERRQ(ierr); 390 } 391 392 ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr); 393 ierr = PetscObjectBaseTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);CHKERRQ(ierr); 394 ierr = PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL);CHKERRQ(ierr); 395 if (isBAIJ) { 396 Mat_MPIBAIJ *baij=(Mat_MPIBAIJ*)mat->data; 397 Mat_SeqBAIJ *spA,*spB; 398 A = baij->A; spA = (Mat_SeqBAIJ*)A->data; A_val = spA->a; 399 B = baij->B; spB = (Mat_SeqBAIJ*)B->data; B_val = spB->a; 400 nz = spA->nz + spB->nz; /* total nonzero entries of mat */ 401 if (!baij->colmap) { 402 ierr = MatCreateColmap_MPIBAIJ_Private(mat);CHKERRQ(ierr); 403 } 404 colmap = baij->colmap; 405 ierr = MatGetColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr); 406 ierr = MatGetColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr); 407 408 if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */ 409 PetscInt *garray; 410 ierr = PetscMalloc1(B->cmap->n,&garray);CHKERRQ(ierr); 411 for (i=0; i<baij->B->cmap->n/bs; i++) { 412 for (j=0; j<bs; j++) { 413 garray[i*bs+j] = bs*baij->garray[i]+j; 414 } 415 } 416 ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,garray,&c->vscale);CHKERRQ(ierr); 417 ierr = VecPinToCPU(c->vscale,PETSC_TRUE);CHKERRQ(ierr); 418 ierr = PetscFree(garray);CHKERRQ(ierr); 419 } 420 } else if (isSELL) { 421 Mat_MPISELL *sell=(Mat_MPISELL*)mat->data; 422 Mat_SeqSELL *spA,*spB; 423 A = sell->A; spA = (Mat_SeqSELL*)A->data; A_val = spA->val; 424 B = sell->B; spB = (Mat_SeqSELL*)B->data; B_val = spB->val; 425 nz = spA->nz + spB->nz; /* total nonzero entries of mat */ 426 if (!sell->colmap) { 427 /* Allow access to data structures of local part of matrix 428 - creates aij->colmap which maps global column number to local number in part B */ 429 ierr = MatCreateColmap_MPISELL_Private(mat);CHKERRQ(ierr); 430 } 431 colmap = sell->colmap; 432 ierr = MatGetColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr); 433 ierr = MatGetColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr); 434 435 bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */ 436 437 if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */ 438 ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,sell->garray,&c->vscale);CHKERRQ(ierr); 439 ierr = VecPinToCPU(c->vscale,PETSC_TRUE);CHKERRQ(ierr); 440 } 441 } else { 442 Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data; 443 Mat_SeqAIJ *spA,*spB; 444 A = aij->A; spA = (Mat_SeqAIJ*)A->data; A_val = spA->a; 445 B = aij->B; spB = (Mat_SeqAIJ*)B->data; B_val = spB->a; 446 nz = spA->nz + spB->nz; /* total nonzero entries of mat */ 447 if (!aij->colmap) { 448 /* Allow access to data structures of local part of matrix 449 - creates aij->colmap which maps global column number to local number in part B */ 450 ierr = MatCreateColmap_MPIAIJ_Private(mat);CHKERRQ(ierr); 451 } 452 colmap = aij->colmap; 453 ierr = MatGetColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr); 454 ierr = MatGetColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr); 455 456 bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */ 457 458 if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */ 459 ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,aij->garray,&c->vscale);CHKERRQ(ierr); 460 ierr = VecPinToCPU(c->vscale,PETSC_TRUE);CHKERRQ(ierr); 461 } 462 } 463 464 m = mat->rmap->n/bs; 465 cstart = mat->cmap->rstart/bs; 466 cend = mat->cmap->rend/bs; 467 468 ierr = PetscMalloc2(nis,&c->ncolumns,nis,&c->columns);CHKERRQ(ierr); 469 ierr = PetscMalloc1(nis,&c->nrows);CHKERRQ(ierr); 470 ierr = PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));CHKERRQ(ierr); 471 472 if (c->htype[0] == 'd') { 473 ierr = PetscMalloc1(nz,&Jentry);CHKERRQ(ierr); 474 ierr = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));CHKERRQ(ierr); 475 c->matentry = Jentry; 476 } else if (c->htype[0] == 'w') { 477 ierr = PetscMalloc1(nz,&Jentry2);CHKERRQ(ierr); 478 ierr = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry2));CHKERRQ(ierr); 479 c->matentry2 = Jentry2; 480 } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"htype is not supported"); 481 482 ierr = PetscMalloc2(m+1,&rowhit,m+1,&valaddrhit);CHKERRQ(ierr); 483 nz = 0; 484 ierr = ISColoringGetIS(iscoloring,PETSC_OWN_POINTER, PETSC_IGNORE,&c->isa);CHKERRQ(ierr); 485 486 if (ctype == IS_COLORING_GLOBAL) { 487 ierr = MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size);CHKERRQ(ierr); 488 ierr = PetscMalloc2(size,&ncolsonproc,size,&disp);CHKERRQ(ierr); 489 } 490 491 for (i=0; i<nis; i++) { /* for each local color */ 492 ierr = ISGetLocalSize(c->isa[i],&n);CHKERRQ(ierr); 493 ierr = ISGetIndices(c->isa[i],&is);CHKERRQ(ierr); 494 495 c->ncolumns[i] = n; /* local number of columns of this color on this process */ 496 c->columns[i] = (PetscInt*)is; 497 498 if (ctype == IS_COLORING_GLOBAL) { 499 /* Determine nctot, the total (parallel) number of columns of this color */ 500 /* ncolsonproc[j]: local ncolumns on proc[j] of this color */ 501 ierr = PetscMPIIntCast(n,&nn);CHKERRQ(ierr); 502 ierr = MPI_Allgather(&nn,1,MPI_INT,ncolsonproc,1,MPI_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr); 503 nctot = 0; for (j=0; j<size; j++) nctot += ncolsonproc[j]; 504 if (!nctot) { 505 ierr = PetscInfo(mat,"Coloring of matrix has some unneeded colors with no corresponding rows\n");CHKERRQ(ierr); 506 } 507 508 disp[0] = 0; 509 for (j=1; j<size; j++) { 510 disp[j] = disp[j-1] + ncolsonproc[j-1]; 511 } 512 513 /* Get cols, the complete list of columns for this color on each process */ 514 ierr = PetscMalloc1(nctot+1,&cols);CHKERRQ(ierr); 515 ierr = MPI_Allgatherv((void*)is,n,MPIU_INT,cols,ncolsonproc,disp,MPIU_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr); 516 } else if (ctype == IS_COLORING_LOCAL) { 517 /* Determine local number of columns of this color on this process, including ghost points */ 518 nctot = n; 519 cols = (PetscInt*)is; 520 } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not provided for this MatFDColoring type"); 521 522 /* Mark all rows affect by these columns */ 523 ierr = PetscArrayzero(rowhit,m);CHKERRQ(ierr); 524 bs2 = bs*bs; 525 nrows_i = 0; 526 for (j=0; j<nctot; j++) { /* loop over columns*/ 527 if (ctype == IS_COLORING_LOCAL) { 528 col = ltog[cols[j]]; 529 } else { 530 col = cols[j]; 531 } 532 if (col >= cstart && col < cend) { /* column is in A, diagonal block of mat */ 533 tmp = A_ci[col-cstart]; 534 row = A_cj + tmp; 535 nrows = A_ci[col-cstart+1] - tmp; 536 nrows_i += nrows; 537 538 /* loop over columns of A marking them in rowhit */ 539 for (k=0; k<nrows; k++) { 540 /* set valaddrhit for part A */ 541 spidx = bs2*spidxA[tmp + k]; 542 valaddrhit[*row] = &A_val[spidx]; 543 rowhit[*row++] = col - cstart + 1; /* local column index */ 544 } 545 } else { /* column is in B, off-diagonal block of mat */ 546 #if defined(PETSC_USE_CTABLE) 547 ierr = PetscTableFind(colmap,col+1,&colb);CHKERRQ(ierr); 548 colb--; 549 #else 550 colb = colmap[col] - 1; /* local column index */ 551 #endif 552 if (colb == -1) { 553 nrows = 0; 554 } else { 555 colb = colb/bs; 556 tmp = B_ci[colb]; 557 row = B_cj + tmp; 558 nrows = B_ci[colb+1] - tmp; 559 } 560 nrows_i += nrows; 561 /* loop over columns of B marking them in rowhit */ 562 for (k=0; k<nrows; k++) { 563 /* set valaddrhit for part B */ 564 spidx = bs2*spidxB[tmp + k]; 565 valaddrhit[*row] = &B_val[spidx]; 566 rowhit[*row++] = colb + 1 + cend - cstart; /* local column index */ 567 } 568 } 569 } 570 c->nrows[i] = nrows_i; 571 572 if (c->htype[0] == 'd') { 573 for (j=0; j<m; j++) { 574 if (rowhit[j]) { 575 Jentry[nz].row = j; /* local row index */ 576 Jentry[nz].col = rowhit[j] - 1; /* local column index */ 577 Jentry[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */ 578 nz++; 579 } 580 } 581 } else { /* c->htype == 'wp' */ 582 for (j=0; j<m; j++) { 583 if (rowhit[j]) { 584 Jentry2[nz].row = j; /* local row index */ 585 Jentry2[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */ 586 nz++; 587 } 588 } 589 } 590 if (ctype == IS_COLORING_GLOBAL) { 591 ierr = PetscFree(cols);CHKERRQ(ierr); 592 } 593 } 594 if (ctype == IS_COLORING_GLOBAL) { 595 ierr = PetscFree2(ncolsonproc,disp);CHKERRQ(ierr); 596 } 597 598 if (bcols > 1) { /* reorder Jentry for faster MatFDColoringApply() */ 599 ierr = MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz);CHKERRQ(ierr); 600 } 601 602 if (isBAIJ) { 603 ierr = MatRestoreColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr); 604 ierr = MatRestoreColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr); 605 ierr = PetscMalloc1(bs*mat->rmap->n,&c->dy);CHKERRQ(ierr); 606 } else if (isSELL) { 607 ierr = MatRestoreColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr); 608 ierr = MatRestoreColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr); 609 }else { 610 ierr = MatRestoreColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr); 611 ierr = MatRestoreColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr); 612 } 613 614 ierr = ISColoringRestoreIS(iscoloring,PETSC_OWN_POINTER,&c->isa);CHKERRQ(ierr); 615 ierr = PetscFree2(rowhit,valaddrhit);CHKERRQ(ierr); 616 617 if (ctype == IS_COLORING_LOCAL) { 618 ierr = ISLocalToGlobalMappingRestoreIndices(map,<og);CHKERRQ(ierr); 619 } 620 ierr = PetscInfo3(c,"ncolors %D, brows %D and bcols %D are used.\n",c->ncolors,c->brows,c->bcols);CHKERRQ(ierr); 621 PetscFunctionReturn(0); 622 } 623 624 PetscErrorCode MatFDColoringCreate_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c) 625 { 626 PetscErrorCode ierr; 627 PetscInt bs,nis=iscoloring->n,m=mat->rmap->n; 628 PetscBool isBAIJ,isSELL; 629 630 PetscFunctionBegin; 631 /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian; 632 bcols is chosen s.t. dy-array takes 50% of memory space as mat */ 633 ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr); 634 ierr = PetscObjectBaseTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);CHKERRQ(ierr); 635 ierr = PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL);CHKERRQ(ierr); 636 if (isBAIJ || m == 0) { 637 c->brows = m; 638 c->bcols = 1; 639 } else if (isSELL) { 640 /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */ 641 Mat_MPISELL *sell=(Mat_MPISELL*)mat->data; 642 Mat_SeqSELL *spA,*spB; 643 Mat A,B; 644 PetscInt nz,brows,bcols; 645 PetscReal mem; 646 647 bs = 1; /* only bs=1 is supported for MPISELL matrix */ 648 649 A = sell->A; spA = (Mat_SeqSELL*)A->data; 650 B = sell->B; spB = (Mat_SeqSELL*)B->data; 651 nz = spA->nz + spB->nz; /* total local nonzero entries of mat */ 652 mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt); 653 bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar))); 654 brows = 1000/bcols; 655 if (bcols > nis) bcols = nis; 656 if (brows == 0 || brows > m) brows = m; 657 c->brows = brows; 658 c->bcols = bcols; 659 } else { /* mpiaij matrix */ 660 /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */ 661 Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data; 662 Mat_SeqAIJ *spA,*spB; 663 Mat A,B; 664 PetscInt nz,brows,bcols; 665 PetscReal mem; 666 667 bs = 1; /* only bs=1 is supported for MPIAIJ matrix */ 668 669 A = aij->A; spA = (Mat_SeqAIJ*)A->data; 670 B = aij->B; spB = (Mat_SeqAIJ*)B->data; 671 nz = spA->nz + spB->nz; /* total local nonzero entries of mat */ 672 mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt); 673 bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar))); 674 brows = 1000/bcols; 675 if (bcols > nis) bcols = nis; 676 if (brows == 0 || brows > m) brows = m; 677 c->brows = brows; 678 c->bcols = bcols; 679 } 680 681 c->M = mat->rmap->N/bs; /* set the global rows and columns and local rows */ 682 c->N = mat->cmap->N/bs; 683 c->m = mat->rmap->n/bs; 684 c->rstart = mat->rmap->rstart/bs; 685 c->ncolors = nis; 686 PetscFunctionReturn(0); 687 } 688