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 PetscInt k,cstart,cend,l,row,col,nz,spidx,i,j; 10 PetscScalar dx=0.0,*w3_array,*dy_i,*dy=coloring->dy; 11 PetscScalar *vscale_array; 12 const PetscScalar *xx; 13 PetscReal epsilon=coloring->error_rel,umin=coloring->umin,unorm; 14 Vec w1=coloring->w1,w2=coloring->w2,w3,vscale=coloring->vscale; 15 void *fctx=coloring->fctx; 16 PetscInt ctype=coloring->ctype,nxloc,nrows_k; 17 PetscScalar *valaddr; 18 MatEntry *Jentry=coloring->matentry; 19 MatEntry2 *Jentry2=coloring->matentry2; 20 const PetscInt ncolors=coloring->ncolors,*ncolumns=coloring->ncolumns,*nrows=coloring->nrows; 21 PetscInt bs=J->rmap->bs; 22 23 PetscFunctionBegin; 24 PetscCall(VecBindToCPU(x1,PETSC_TRUE)); 25 /* (1) Set w1 = F(x1) */ 26 if (!coloring->fset) { 27 PetscCall(PetscLogEventBegin(MAT_FDColoringFunction,coloring,0,0,0)); 28 PetscCall((*f)(sctx,x1,w1,fctx)); 29 PetscCall(PetscLogEventEnd(MAT_FDColoringFunction,coloring,0,0,0)); 30 } else { 31 coloring->fset = PETSC_FALSE; 32 } 33 34 /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */ 35 PetscCall(VecGetLocalSize(x1,&nxloc)); 36 if (coloring->htype[0] == 'w') { 37 /* vscale = dx is a constant scalar */ 38 PetscCall(VecNorm(x1,NORM_2,&unorm)); 39 dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon); 40 } else { 41 PetscCall(VecGetArrayRead(x1,&xx)); 42 PetscCall(VecGetArray(vscale,&vscale_array)); 43 for (col=0; col<nxloc; col++) { 44 dx = xx[col]; 45 if (PetscAbsScalar(dx) < umin) { 46 if (PetscRealPart(dx) >= 0.0) dx = umin; 47 else if (PetscRealPart(dx) < 0.0) dx = -umin; 48 } 49 dx *= epsilon; 50 vscale_array[col] = 1.0/dx; 51 } 52 PetscCall(VecRestoreArrayRead(x1,&xx)); 53 PetscCall(VecRestoreArray(vscale,&vscale_array)); 54 } 55 if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') { 56 PetscCall(VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD)); 57 PetscCall(VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD)); 58 } 59 60 /* (3) Loop over each color */ 61 if (!coloring->w3) { 62 PetscCall(VecDuplicate(x1,&coloring->w3)); 63 /* Vec is used intensively in particular piece of scalar CPU code; won't benefit from bouncing back and forth to the GPU */ 64 PetscCall(VecBindToCPU(coloring->w3,PETSC_TRUE)); 65 PetscCall(PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3)); 66 } 67 w3 = coloring->w3; 68 69 PetscCall(VecGetOwnershipRange(x1,&cstart,&cend)); /* used by ghosted vscale */ 70 if (vscale) { 71 PetscCall(VecGetArray(vscale,&vscale_array)); 72 } 73 nz = 0; 74 for (k=0; k<ncolors; k++) { 75 coloring->currentcolor = k; 76 77 /* 78 (3-1) Loop over each column associated with color 79 adding the perturbation to the vector w3 = x1 + dx. 80 */ 81 PetscCall(VecCopy(x1,w3)); 82 dy_i = dy; 83 for (i=0; i<bs; i++) { /* Loop over a block of columns */ 84 PetscCall(VecGetArray(w3,&w3_array)); 85 if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */ 86 if (coloring->htype[0] == 'w') { 87 for (l=0; l<ncolumns[k]; l++) { 88 col = i + bs*coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */ 89 w3_array[col] += 1.0/dx; 90 if (i) w3_array[col-1] -= 1.0/dx; /* resume original w3[col-1] */ 91 } 92 } else { /* htype == 'ds' */ 93 vscale_array -= cstart; /* shift pointer so global index can be used */ 94 for (l=0; l<ncolumns[k]; l++) { 95 col = i + bs*coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */ 96 w3_array[col] += 1.0/vscale_array[col]; 97 if (i) w3_array[col-1] -= 1.0/vscale_array[col-1]; /* resume original w3[col-1] */ 98 } 99 vscale_array += cstart; 100 } 101 if (ctype == IS_COLORING_GLOBAL) w3_array += cstart; 102 PetscCall(VecRestoreArray(w3,&w3_array)); 103 104 /* 105 (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations) 106 w2 = F(x1 + dx) - F(x1) 107 */ 108 PetscCall(PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0)); 109 PetscCall(VecPlaceArray(w2,dy_i)); /* place w2 to the array dy_i */ 110 PetscCall((*f)(sctx,w3,w2,fctx)); 111 PetscCall(PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0)); 112 PetscCall(VecAXPY(w2,-1.0,w1)); 113 PetscCall(VecResetArray(w2)); 114 dy_i += nxloc; /* points to dy+i*nxloc */ 115 } 116 117 /* 118 (3-3) Loop over rows of vector, putting results into Jacobian matrix 119 */ 120 nrows_k = nrows[k]; 121 if (coloring->htype[0] == 'w') { 122 for (l=0; l<nrows_k; l++) { 123 row = bs*Jentry2[nz].row; /* local row index */ 124 valaddr = Jentry2[nz++].valaddr; 125 spidx = 0; 126 dy_i = dy; 127 for (i=0; i<bs; i++) { /* column of the block */ 128 for (j=0; j<bs; j++) { /* row of the block */ 129 valaddr[spidx++] = dy_i[row+j]*dx; 130 } 131 dy_i += nxloc; /* points to dy+i*nxloc */ 132 } 133 } 134 } else { /* htype == 'ds' */ 135 for (l=0; l<nrows_k; l++) { 136 row = bs*Jentry[nz].row; /* local row index */ 137 col = bs*Jentry[nz].col; /* local column index */ 138 valaddr = Jentry[nz++].valaddr; 139 spidx = 0; 140 dy_i = dy; 141 for (i=0; i<bs; i++) { /* column of the block */ 142 for (j=0; j<bs; j++) { /* row of the block */ 143 valaddr[spidx++] = dy_i[row+j]*vscale_array[col+i]; 144 } 145 dy_i += nxloc; /* points to dy+i*nxloc */ 146 } 147 } 148 } 149 } 150 PetscCall(MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY)); 151 PetscCall(MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY)); 152 if (vscale) { 153 PetscCall(VecRestoreArray(vscale,&vscale_array)); 154 } 155 156 coloring->currentcolor = -1; 157 PetscCall(VecBindToCPU(x1,PETSC_FALSE)); 158 PetscFunctionReturn(0); 159 } 160 161 /* this is declared PETSC_EXTERN because it is used by MatFDColoringUseDM() which is in the DM library */ 162 PetscErrorCode MatFDColoringApply_AIJ(Mat J,MatFDColoring coloring,Vec x1,void *sctx) 163 { 164 PetscErrorCode (*f)(void*,Vec,Vec,void*) = (PetscErrorCode (*)(void*,Vec,Vec,void*))coloring->f; 165 PetscInt k,cstart,cend,l,row,col,nz; 166 PetscScalar dx=0.0,*y,*w3_array; 167 const PetscScalar *xx; 168 PetscScalar *vscale_array; 169 PetscReal epsilon=coloring->error_rel,umin=coloring->umin,unorm; 170 Vec w1=coloring->w1,w2=coloring->w2,w3,vscale=coloring->vscale; 171 void *fctx=coloring->fctx; 172 ISColoringType ctype=coloring->ctype; 173 PetscInt nxloc,nrows_k; 174 MatEntry *Jentry=coloring->matentry; 175 MatEntry2 *Jentry2=coloring->matentry2; 176 const PetscInt ncolors=coloring->ncolors,*ncolumns=coloring->ncolumns,*nrows=coloring->nrows; 177 PetscBool alreadyboundtocpu; 178 179 PetscFunctionBegin; 180 PetscCall(VecBoundToCPU(x1,&alreadyboundtocpu)); 181 PetscCall(VecBindToCPU(x1,PETSC_TRUE)); 182 PetscCheckFalse((ctype == IS_COLORING_LOCAL) && (J->ops->fdcoloringapply == MatFDColoringApply_AIJ),PetscObjectComm((PetscObject)J),PETSC_ERR_SUP,"Must call MatColoringUseDM() with IS_COLORING_LOCAL"); 183 /* (1) Set w1 = F(x1) */ 184 if (!coloring->fset) { 185 PetscCall(PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0)); 186 PetscCall((*f)(sctx,x1,w1,fctx)); 187 PetscCall(PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0)); 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 PetscCall(VecNorm(x1,NORM_2,&unorm)); 196 dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon); 197 } else { 198 PetscCall(VecGetLocalSize(x1,&nxloc)); 199 PetscCall(VecGetArrayRead(x1,&xx)); 200 PetscCall(VecGetArray(vscale,&vscale_array)); 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 PetscCall(VecRestoreArrayRead(x1,&xx)); 211 PetscCall(VecRestoreArray(vscale,&vscale_array)); 212 } 213 if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') { 214 PetscCall(VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD)); 215 PetscCall(VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD)); 216 } 217 218 /* (3) Loop over each color */ 219 if (!coloring->w3) { 220 PetscCall(VecDuplicate(x1,&coloring->w3)); 221 PetscCall(PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3)); 222 } 223 w3 = coloring->w3; 224 225 PetscCall(VecGetOwnershipRange(x1,&cstart,&cend)); /* used by ghosted vscale */ 226 if (vscale) { 227 PetscCall(VecGetArray(vscale,&vscale_array)); 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 PetscCall(VecCopy(x1,w3)); 249 PetscCall(VecGetArray(w3,&w3_array)); 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 PetscCall(VecRestoreArray(w3,&w3_array)); 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 PetscCall(PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0)); 272 PetscCall(VecPlaceArray(w2,dy_k)); /* place w2 to the array dy_i */ 273 PetscCall((*f)(sctx,w3,w2,fctx)); 274 PetscCall(PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0)); 275 PetscCall(VecAXPY(w2,-1.0,w1)); 276 PetscCall(VecResetArray(w2)); 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 /* The 'useless' ifdef is due to a bug in NVIDIA nvc 21.11, which triggers a segfault on this line. We write it in 289 another way, and it seems work. See https://lists.mcs.anl.gov/pipermail/petsc-users/2021-December/045158.html 290 */ 291 #if defined(PETSC_USE_COMPLEX) 292 PetscScalar *tmp = Jentry2[nz].valaddr; 293 *tmp = dy[row]*dx; 294 #else 295 *(Jentry2[nz].valaddr) = dy[row]*dx; 296 #endif 297 nz++; 298 } 299 } else { /* htype == 'ds' */ 300 for (l=0; l<nrows_k; l++) { 301 row = Jentry[nz].row; /* local row index */ 302 #if defined(PETSC_USE_COMPLEX) /* See https://lists.mcs.anl.gov/pipermail/petsc-users/2021-December/045158.html */ 303 PetscScalar *tmp = Jentry[nz].valaddr; 304 *tmp = dy[row]*vscale_array[Jentry[nz].col]; 305 #else 306 *(Jentry[nz].valaddr) = dy[row]*vscale_array[Jentry[nz].col]; 307 #endif 308 nz++; 309 } 310 } 311 } 312 } else { /* bcols == 1 */ 313 for (k=0; k<ncolors; k++) { 314 coloring->currentcolor = k; 315 316 /* 317 (3-1) Loop over each column associated with color 318 adding the perturbation to the vector w3 = x1 + dx. 319 */ 320 PetscCall(VecCopy(x1,w3)); 321 PetscCall(VecGetArray(w3,&w3_array)); 322 if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */ 323 if (coloring->htype[0] == 'w') { 324 for (l=0; l<ncolumns[k]; l++) { 325 col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */ 326 w3_array[col] += 1.0/dx; 327 } 328 } else { /* htype == 'ds' */ 329 vscale_array -= cstart; /* shift pointer so global index can be used */ 330 for (l=0; l<ncolumns[k]; l++) { 331 col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */ 332 w3_array[col] += 1.0/vscale_array[col]; 333 } 334 vscale_array += cstart; 335 } 336 if (ctype == IS_COLORING_GLOBAL) w3_array += cstart; 337 PetscCall(VecRestoreArray(w3,&w3_array)); 338 339 /* 340 (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations) 341 w2 = F(x1 + dx) - F(x1) 342 */ 343 PetscCall(PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0)); 344 PetscCall((*f)(sctx,w3,w2,fctx)); 345 PetscCall(PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0)); 346 PetscCall(VecAXPY(w2,-1.0,w1)); 347 348 /* 349 (3-3) Loop over rows of vector, putting results into Jacobian matrix 350 */ 351 nrows_k = nrows[k]; 352 PetscCall(VecGetArray(w2,&y)); 353 if (coloring->htype[0] == 'w') { 354 for (l=0; l<nrows_k; l++) { 355 row = Jentry2[nz].row; /* local row index */ 356 #if defined(PETSC_USE_COMPLEX) /* See https://lists.mcs.anl.gov/pipermail/petsc-users/2021-December/045158.html */ 357 PetscScalar *tmp = Jentry2[nz].valaddr; 358 *tmp = y[row]*dx; 359 #else 360 *(Jentry2[nz].valaddr) = y[row]*dx; 361 #endif 362 nz++; 363 } 364 } else { /* htype == 'ds' */ 365 for (l=0; l<nrows_k; l++) { 366 row = Jentry[nz].row; /* local row index */ 367 #if defined(PETSC_USE_COMPLEX) /* See https://lists.mcs.anl.gov/pipermail/petsc-users/2021-December/045158.html */ 368 PetscScalar *tmp = Jentry[nz].valaddr; 369 *tmp = y[row]*vscale_array[Jentry[nz].col]; 370 #else 371 *(Jentry[nz].valaddr) = y[row]*vscale_array[Jentry[nz].col]; 372 #endif 373 nz++; 374 } 375 } 376 PetscCall(VecRestoreArray(w2,&y)); 377 } 378 } 379 380 #if defined(PETSC_HAVE_VIENNACL) || defined(PETSC_HAVE_CUDA) 381 if (J->offloadmask != PETSC_OFFLOAD_UNALLOCATED) J->offloadmask = PETSC_OFFLOAD_CPU; 382 #endif 383 PetscCall(MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY)); 384 PetscCall(MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY)); 385 if (vscale) { 386 PetscCall(VecRestoreArray(vscale,&vscale_array)); 387 } 388 coloring->currentcolor = -1; 389 if (!alreadyboundtocpu) PetscCall(VecBindToCPU(x1,PETSC_FALSE)); 390 PetscFunctionReturn(0); 391 } 392 393 PetscErrorCode MatFDColoringSetUp_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c) 394 { 395 PetscMPIInt size,*ncolsonproc,*disp,nn; 396 PetscInt i,n,nrows,nrows_i,j,k,m,ncols,col,*rowhit,cstart,cend,colb; 397 const PetscInt *is,*A_ci,*A_cj,*B_ci,*B_cj,*row=NULL,*ltog=NULL; 398 PetscInt nis=iscoloring->n,nctot,*cols,tmp = 0; 399 ISLocalToGlobalMapping map=mat->cmap->mapping; 400 PetscInt ctype=c->ctype,*spidxA,*spidxB,nz,bs,bs2,spidx; 401 Mat A,B; 402 PetscScalar *A_val,*B_val,**valaddrhit; 403 MatEntry *Jentry; 404 MatEntry2 *Jentry2; 405 PetscBool isBAIJ,isSELL; 406 PetscInt bcols=c->bcols; 407 #if defined(PETSC_USE_CTABLE) 408 PetscTable colmap=NULL; 409 #else 410 PetscInt *colmap=NULL; /* local col number of off-diag col */ 411 #endif 412 413 PetscFunctionBegin; 414 if (ctype == IS_COLORING_LOCAL) { 415 PetscCheck(map,PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_INCOMP,"When using ghosted differencing matrix must have local to global mapping provided with MatSetLocalToGlobalMapping"); 416 PetscCall(ISLocalToGlobalMappingGetIndices(map,<og)); 417 } 418 419 PetscCall(MatGetBlockSize(mat,&bs)); 420 PetscCall(PetscObjectBaseTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ)); 421 PetscCall(PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL)); 422 if (isBAIJ) { 423 Mat_MPIBAIJ *baij=(Mat_MPIBAIJ*)mat->data; 424 Mat_SeqBAIJ *spA,*spB; 425 A = baij->A; spA = (Mat_SeqBAIJ*)A->data; A_val = spA->a; 426 B = baij->B; spB = (Mat_SeqBAIJ*)B->data; B_val = spB->a; 427 nz = spA->nz + spB->nz; /* total nonzero entries of mat */ 428 if (!baij->colmap) { 429 PetscCall(MatCreateColmap_MPIBAIJ_Private(mat)); 430 } 431 colmap = baij->colmap; 432 PetscCall(MatGetColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL)); 433 PetscCall(MatGetColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL)); 434 435 if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */ 436 PetscInt *garray; 437 PetscCall(PetscMalloc1(B->cmap->n,&garray)); 438 for (i=0; i<baij->B->cmap->n/bs; i++) { 439 for (j=0; j<bs; j++) { 440 garray[i*bs+j] = bs*baij->garray[i]+j; 441 } 442 } 443 PetscCall(VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,garray,&c->vscale)); 444 PetscCall(VecBindToCPU(c->vscale,PETSC_TRUE)); 445 PetscCall(PetscFree(garray)); 446 } 447 } else if (isSELL) { 448 Mat_MPISELL *sell=(Mat_MPISELL*)mat->data; 449 Mat_SeqSELL *spA,*spB; 450 A = sell->A; spA = (Mat_SeqSELL*)A->data; A_val = spA->val; 451 B = sell->B; spB = (Mat_SeqSELL*)B->data; B_val = spB->val; 452 nz = spA->nz + spB->nz; /* total nonzero entries of mat */ 453 if (!sell->colmap) { 454 /* Allow access to data structures of local part of matrix 455 - creates aij->colmap which maps global column number to local number in part B */ 456 PetscCall(MatCreateColmap_MPISELL_Private(mat)); 457 } 458 colmap = sell->colmap; 459 PetscCall(MatGetColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL)); 460 PetscCall(MatGetColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL)); 461 462 bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */ 463 464 if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */ 465 PetscCall(VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,sell->garray,&c->vscale)); 466 PetscCall(VecBindToCPU(c->vscale,PETSC_TRUE)); 467 } 468 } else { 469 Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data; 470 Mat_SeqAIJ *spA,*spB; 471 A = aij->A; spA = (Mat_SeqAIJ*)A->data; A_val = spA->a; 472 B = aij->B; spB = (Mat_SeqAIJ*)B->data; B_val = spB->a; 473 nz = spA->nz + spB->nz; /* total nonzero entries of mat */ 474 if (!aij->colmap) { 475 /* Allow access to data structures of local part of matrix 476 - creates aij->colmap which maps global column number to local number in part B */ 477 PetscCall(MatCreateColmap_MPIAIJ_Private(mat)); 478 } 479 colmap = aij->colmap; 480 PetscCall(MatGetColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL)); 481 PetscCall(MatGetColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL)); 482 483 bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */ 484 485 if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */ 486 PetscCall(VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,aij->garray,&c->vscale)); 487 PetscCall(VecBindToCPU(c->vscale,PETSC_TRUE)); 488 } 489 } 490 491 m = mat->rmap->n/bs; 492 cstart = mat->cmap->rstart/bs; 493 cend = mat->cmap->rend/bs; 494 495 PetscCall(PetscMalloc2(nis,&c->ncolumns,nis,&c->columns)); 496 PetscCall(PetscMalloc1(nis,&c->nrows)); 497 PetscCall(PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt))); 498 499 if (c->htype[0] == 'd') { 500 PetscCall(PetscMalloc1(nz,&Jentry)); 501 PetscCall(PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry))); 502 c->matentry = Jentry; 503 } else if (c->htype[0] == 'w') { 504 PetscCall(PetscMalloc1(nz,&Jentry2)); 505 PetscCall(PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry2))); 506 c->matentry2 = Jentry2; 507 } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"htype is not supported"); 508 509 PetscCall(PetscMalloc2(m+1,&rowhit,m+1,&valaddrhit)); 510 nz = 0; 511 PetscCall(ISColoringGetIS(iscoloring,PETSC_OWN_POINTER, PETSC_IGNORE,&c->isa)); 512 513 if (ctype == IS_COLORING_GLOBAL) { 514 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size)); 515 PetscCall(PetscMalloc2(size,&ncolsonproc,size,&disp)); 516 } 517 518 for (i=0; i<nis; i++) { /* for each local color */ 519 PetscCall(ISGetLocalSize(c->isa[i],&n)); 520 PetscCall(ISGetIndices(c->isa[i],&is)); 521 522 c->ncolumns[i] = n; /* local number of columns of this color on this process */ 523 c->columns[i] = (PetscInt*)is; 524 525 if (ctype == IS_COLORING_GLOBAL) { 526 /* Determine nctot, the total (parallel) number of columns of this color */ 527 /* ncolsonproc[j]: local ncolumns on proc[j] of this color */ 528 PetscCall(PetscMPIIntCast(n,&nn)); 529 PetscCallMPI(MPI_Allgather(&nn,1,MPI_INT,ncolsonproc,1,MPI_INT,PetscObjectComm((PetscObject)mat))); 530 nctot = 0; for (j=0; j<size; j++) nctot += ncolsonproc[j]; 531 if (!nctot) { 532 PetscCall(PetscInfo(mat,"Coloring of matrix has some unneeded colors with no corresponding rows\n")); 533 } 534 535 disp[0] = 0; 536 for (j=1; j<size; j++) { 537 disp[j] = disp[j-1] + ncolsonproc[j-1]; 538 } 539 540 /* Get cols, the complete list of columns for this color on each process */ 541 PetscCall(PetscMalloc1(nctot+1,&cols)); 542 PetscCallMPI(MPI_Allgatherv((void*)is,n,MPIU_INT,cols,ncolsonproc,disp,MPIU_INT,PetscObjectComm((PetscObject)mat))); 543 } else if (ctype == IS_COLORING_LOCAL) { 544 /* Determine local number of columns of this color on this process, including ghost points */ 545 nctot = n; 546 cols = (PetscInt*)is; 547 } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not provided for this MatFDColoring type"); 548 549 /* Mark all rows affect by these columns */ 550 PetscCall(PetscArrayzero(rowhit,m)); 551 bs2 = bs*bs; 552 nrows_i = 0; 553 for (j=0; j<nctot; j++) { /* loop over columns*/ 554 if (ctype == IS_COLORING_LOCAL) { 555 col = ltog[cols[j]]; 556 } else { 557 col = cols[j]; 558 } 559 if (col >= cstart && col < cend) { /* column is in A, diagonal block of mat */ 560 tmp = A_ci[col-cstart]; 561 row = A_cj + tmp; 562 nrows = A_ci[col-cstart+1] - tmp; 563 nrows_i += nrows; 564 565 /* loop over columns of A marking them in rowhit */ 566 for (k=0; k<nrows; k++) { 567 /* set valaddrhit for part A */ 568 spidx = bs2*spidxA[tmp + k]; 569 valaddrhit[*row] = &A_val[spidx]; 570 rowhit[*row++] = col - cstart + 1; /* local column index */ 571 } 572 } else { /* column is in B, off-diagonal block of mat */ 573 #if defined(PETSC_USE_CTABLE) 574 PetscCall(PetscTableFind(colmap,col+1,&colb)); 575 colb--; 576 #else 577 colb = colmap[col] - 1; /* local column index */ 578 #endif 579 if (colb == -1) { 580 nrows = 0; 581 } else { 582 colb = colb/bs; 583 tmp = B_ci[colb]; 584 row = B_cj + tmp; 585 nrows = B_ci[colb+1] - tmp; 586 } 587 nrows_i += nrows; 588 /* loop over columns of B marking them in rowhit */ 589 for (k=0; k<nrows; k++) { 590 /* set valaddrhit for part B */ 591 spidx = bs2*spidxB[tmp + k]; 592 valaddrhit[*row] = &B_val[spidx]; 593 rowhit[*row++] = colb + 1 + cend - cstart; /* local column index */ 594 } 595 } 596 } 597 c->nrows[i] = nrows_i; 598 599 if (c->htype[0] == 'd') { 600 for (j=0; j<m; j++) { 601 if (rowhit[j]) { 602 Jentry[nz].row = j; /* local row index */ 603 Jentry[nz].col = rowhit[j] - 1; /* local column index */ 604 Jentry[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */ 605 nz++; 606 } 607 } 608 } else { /* c->htype == 'wp' */ 609 for (j=0; j<m; j++) { 610 if (rowhit[j]) { 611 Jentry2[nz].row = j; /* local row index */ 612 Jentry2[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */ 613 nz++; 614 } 615 } 616 } 617 if (ctype == IS_COLORING_GLOBAL) { 618 PetscCall(PetscFree(cols)); 619 } 620 } 621 if (ctype == IS_COLORING_GLOBAL) { 622 PetscCall(PetscFree2(ncolsonproc,disp)); 623 } 624 625 if (bcols > 1) { /* reorder Jentry for faster MatFDColoringApply() */ 626 PetscCall(MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz)); 627 } 628 629 if (isBAIJ) { 630 PetscCall(MatRestoreColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL)); 631 PetscCall(MatRestoreColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL)); 632 PetscCall(PetscMalloc1(bs*mat->rmap->n,&c->dy)); 633 } else if (isSELL) { 634 PetscCall(MatRestoreColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL)); 635 PetscCall(MatRestoreColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL)); 636 } else { 637 PetscCall(MatRestoreColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL)); 638 PetscCall(MatRestoreColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL)); 639 } 640 641 PetscCall(ISColoringRestoreIS(iscoloring,PETSC_OWN_POINTER,&c->isa)); 642 PetscCall(PetscFree2(rowhit,valaddrhit)); 643 644 if (ctype == IS_COLORING_LOCAL) { 645 PetscCall(ISLocalToGlobalMappingRestoreIndices(map,<og)); 646 } 647 PetscCall(PetscInfo(c,"ncolors %" PetscInt_FMT ", brows %" PetscInt_FMT " and bcols %" PetscInt_FMT " are used.\n",c->ncolors,c->brows,c->bcols)); 648 PetscFunctionReturn(0); 649 } 650 651 PetscErrorCode MatFDColoringCreate_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c) 652 { 653 PetscInt bs,nis=iscoloring->n,m=mat->rmap->n; 654 PetscBool isBAIJ,isSELL; 655 656 PetscFunctionBegin; 657 /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian; 658 bcols is chosen s.t. dy-array takes 50% of memory space as mat */ 659 PetscCall(MatGetBlockSize(mat,&bs)); 660 PetscCall(PetscObjectBaseTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ)); 661 PetscCall(PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL)); 662 if (isBAIJ || m == 0) { 663 c->brows = m; 664 c->bcols = 1; 665 } else if (isSELL) { 666 /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */ 667 Mat_MPISELL *sell=(Mat_MPISELL*)mat->data; 668 Mat_SeqSELL *spA,*spB; 669 Mat A,B; 670 PetscInt nz,brows,bcols; 671 PetscReal mem; 672 673 bs = 1; /* only bs=1 is supported for MPISELL matrix */ 674 675 A = sell->A; spA = (Mat_SeqSELL*)A->data; 676 B = sell->B; spB = (Mat_SeqSELL*)B->data; 677 nz = spA->nz + spB->nz; /* total local nonzero entries of mat */ 678 mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt); 679 bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar))); 680 brows = 1000/bcols; 681 if (bcols > nis) bcols = nis; 682 if (brows == 0 || brows > m) brows = m; 683 c->brows = brows; 684 c->bcols = bcols; 685 } else { /* mpiaij matrix */ 686 /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */ 687 Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data; 688 Mat_SeqAIJ *spA,*spB; 689 Mat A,B; 690 PetscInt nz,brows,bcols; 691 PetscReal mem; 692 693 bs = 1; /* only bs=1 is supported for MPIAIJ matrix */ 694 695 A = aij->A; spA = (Mat_SeqAIJ*)A->data; 696 B = aij->B; spB = (Mat_SeqAIJ*)B->data; 697 nz = spA->nz + spB->nz; /* total local nonzero entries of mat */ 698 mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt); 699 bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar))); 700 brows = 1000/bcols; 701 if (bcols > nis) bcols = nis; 702 if (brows == 0 || brows > m) brows = m; 703 c->brows = brows; 704 c->bcols = bcols; 705 } 706 707 c->M = mat->rmap->N/bs; /* set the global rows and columns and local rows */ 708 c->N = mat->cmap->N/bs; 709 c->m = mat->rmap->n/bs; 710 c->rstart = mat->rmap->rstart/bs; 711 c->ncolors = nis; 712 PetscFunctionReturn(0); 713 } 714 715 /*@C 716 717 MatFDColoringSetValues - takes a matrix in compressed color format and enters the matrix into a PETSc Mat 718 719 Collective on J 720 721 Input Parameters: 722 + J - the sparse matrix 723 . coloring - created with MatFDColoringCreate() and a local coloring 724 - y - column major storage of matrix values with one color of values per column, the number of rows of y should match 725 the number of local rows of J and the number of columns is the number of colors. 726 727 Level: intermediate 728 729 Notes: the matrix in compressed color format may come from an Automatic Differentiation code 730 731 The code will be slightly faster if MatFDColoringSetBlockSize(coloring,PETSC_DEFAULT,nc); is called immediately after creating the coloring 732 733 .seealso: MatFDColoringCreate(), ISColoring, ISColoringCreate(), ISColoringSetType(), IS_COLORING_LOCAL, MatFDColoringSetBlockSize() 734 735 @*/ 736 PetscErrorCode MatFDColoringSetValues(Mat J,MatFDColoring coloring,const PetscScalar *y) 737 { 738 MatEntry2 *Jentry2; 739 PetscInt row,i,nrows_k,l,ncolors,nz = 0,bcols,nbcols = 0; 740 const PetscInt *nrows; 741 PetscBool eq; 742 743 PetscFunctionBegin; 744 PetscValidHeaderSpecific(J,MAT_CLASSID,1); 745 PetscValidHeaderSpecific(coloring,MAT_FDCOLORING_CLASSID,2); 746 PetscCall(PetscObjectCompareId((PetscObject)J,coloring->matid,&eq)); 747 PetscCheck(eq,PetscObjectComm((PetscObject)J),PETSC_ERR_ARG_WRONG,"Matrix used with MatFDColoringSetValues() must be that used with MatFDColoringCreate()"); 748 Jentry2 = coloring->matentry2; 749 nrows = coloring->nrows; 750 ncolors = coloring->ncolors; 751 bcols = coloring->bcols; 752 753 for (i=0; i<ncolors; i+=bcols) { 754 nrows_k = nrows[nbcols++]; 755 for (l=0; l<nrows_k; l++) { 756 row = Jentry2[nz].row; /* local row index */ 757 *(Jentry2[nz++].valaddr) = y[row]; 758 } 759 y += bcols*coloring->m; 760 } 761 PetscCall(MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY)); 762 PetscCall(MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY)); 763 PetscFunctionReturn(0); 764 } 765