Actual source code: fdmpiaij.c
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>
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;
25: VecBindToCPU(x1,PETSC_TRUE);
26: /* (1) Set w1 = F(x1) */
27: if (!coloring->fset) {
28: PetscLogEventBegin(MAT_FDColoringFunction,coloring,0,0,0);
29: (*f)(sctx,x1,w1,fctx);
30: PetscLogEventEnd(MAT_FDColoringFunction,coloring,0,0,0);
31: } else {
32: coloring->fset = PETSC_FALSE;
33: }
35: /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */
36: VecGetLocalSize(x1,&nxloc);
37: if (coloring->htype[0] == 'w') {
38: /* vscale = dx is a constant scalar */
39: VecNorm(x1,NORM_2,&unorm);
40: dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon);
41: } else {
42: VecGetArrayRead(x1,&xx);
43: VecGetArray(vscale,&vscale_array);
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: VecRestoreArrayRead(x1,&xx);
54: VecRestoreArray(vscale,&vscale_array);
55: }
56: if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') {
57: VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD);
58: VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD);
59: }
61: /* (3) Loop over each color */
62: if (!coloring->w3) {
63: VecDuplicate(x1,&coloring->w3);
64: /* Vec is used intensively in particular piece of scalar CPU code; won't benefit from bouncing back and forth to the GPU */
65: VecBindToCPU(coloring->w3,PETSC_TRUE);
66: PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);
67: }
68: w3 = coloring->w3;
70: VecGetOwnershipRange(x1,&cstart,&cend); /* used by ghosted vscale */
71: if (vscale) {
72: VecGetArray(vscale,&vscale_array);
73: }
74: nz = 0;
75: for (k=0; k<ncolors; k++) {
76: coloring->currentcolor = k;
78: /*
79: (3-1) Loop over each column associated with color
80: adding the perturbation to the vector w3 = x1 + dx.
81: */
82: VecCopy(x1,w3);
83: dy_i = dy;
84: for (i=0; i<bs; i++) { /* Loop over a block of columns */
85: VecGetArray(w3,&w3_array);
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: VecRestoreArray(w3,&w3_array);
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: PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);
110: VecPlaceArray(w2,dy_i); /* place w2 to the array dy_i */
111: (*f)(sctx,w3,w2,fctx);
112: PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);
113: VecAXPY(w2,-1.0,w1);
114: VecResetArray(w2);
115: dy_i += nxloc; /* points to dy+i*nxloc */
116: }
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: MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);
152: MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);
153: if (vscale) {
154: VecRestoreArray(vscale,&vscale_array);
155: }
157: coloring->currentcolor = -1;
158: VecBindToCPU(x1,PETSC_FALSE);
159: return(0);
160: }
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;
181: VecBindToCPU(x1,PETSC_TRUE);
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: PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);
186: (*f)(sctx,x1,w1,fctx);
187: PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);
188: } else {
189: coloring->fset = PETSC_FALSE;
190: }
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: VecNorm(x1,NORM_2,&unorm);
196: dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon);
197: } else {
198: VecGetLocalSize(x1,&nxloc);
199: VecGetArrayRead(x1,&xx);
200: 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: VecRestoreArrayRead(x1,&xx);
211: VecRestoreArray(vscale,&vscale_array);
212: }
213: if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') {
214: VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD);
215: VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD);
216: }
218: /* (3) Loop over each color */
219: if (!coloring->w3) {
220: VecDuplicate(x1,&coloring->w3);
221: PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);
222: }
223: w3 = coloring->w3;
225: VecGetOwnershipRange(x1,&cstart,&cend); /* used by ghosted vscale */
226: if (vscale) {
227: VecGetArray(vscale,&vscale_array);
228: }
229: nz = 0;
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;
235: nbcols = 0;
236: for (k=0; k<ncolors; k+=bcols) {
238: /*
239: (3-1) Loop over each column associated with color
240: adding the perturbation to the vector w3 = x1 + dx.
241: */
243: dy_k = dy;
244: if (k + bcols > ncolors) bcols = ncolors - k;
245: for (i=0; i<bcols; i++) {
246: coloring->currentcolor = k+i;
248: VecCopy(x1,w3);
249: 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: VecRestoreArray(w3,&w3_array);
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: PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);
272: VecPlaceArray(w2,dy_k); /* place w2 to the array dy_i */
273: (*f)(sctx,w3,w2,fctx);
274: PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);
275: VecAXPY(w2,-1.0,w1);
276: VecResetArray(w2);
277: dy_k += m; /* points to dy+i*nxloc */
278: }
280: /*
281: (3-3) Loop over block rows of vector, putting results into Jacobian matrix
282: */
283: nrows_k = nrows[nbcols++];
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;
316: /*
317: (3-1) Loop over each column associated with color
318: adding the perturbation to the vector w3 = x1 + dx.
319: */
320: VecCopy(x1,w3);
321: 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: VecRestoreArray(w3,&w3_array);
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: PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);
344: (*f)(sctx,w3,w2,fctx);
345: PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);
346: VecAXPY(w2,-1.0,w1);
348: /*
349: (3-3) Loop over rows of vector, putting results into Jacobian matrix
350: */
351: nrows_k = nrows[k];
352: 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: VecRestoreArray(w2,&y);
377: }
378: }
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: MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);
384: MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);
385: if (vscale) {
386: VecRestoreArray(vscale,&vscale_array);
387: }
388: coloring->currentcolor = -1;
389: VecBindToCPU(x1,PETSC_FALSE);
390: return(0);
391: }
393: PetscErrorCode MatFDColoringSetUp_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
394: {
395: PetscErrorCode ierr;
396: PetscMPIInt size,*ncolsonproc,*disp,nn;
397: PetscInt i,n,nrows,nrows_i,j,k,m,ncols,col,*rowhit,cstart,cend,colb;
398: const PetscInt *is,*A_ci,*A_cj,*B_ci,*B_cj,*row=NULL,*ltog=NULL;
399: PetscInt nis=iscoloring->n,nctot,*cols,tmp = 0;
400: ISLocalToGlobalMapping map=mat->cmap->mapping;
401: PetscInt ctype=c->ctype,*spidxA,*spidxB,nz,bs,bs2,spidx;
402: Mat A,B;
403: PetscScalar *A_val,*B_val,**valaddrhit;
404: MatEntry *Jentry;
405: MatEntry2 *Jentry2;
406: PetscBool isBAIJ,isSELL;
407: PetscInt bcols=c->bcols;
408: #if defined(PETSC_USE_CTABLE)
409: PetscTable colmap=NULL;
410: #else
411: PetscInt *colmap=NULL; /* local col number of off-diag col */
412: #endif
415: if (ctype == IS_COLORING_LOCAL) {
416: if (!map) SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_INCOMP,"When using ghosted differencing matrix must have local to global mapping provided with MatSetLocalToGlobalMapping");
417: ISLocalToGlobalMappingGetIndices(map,<og);
418: }
420: MatGetBlockSize(mat,&bs);
421: PetscObjectBaseTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);
422: PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL);
423: if (isBAIJ) {
424: Mat_MPIBAIJ *baij=(Mat_MPIBAIJ*)mat->data;
425: Mat_SeqBAIJ *spA,*spB;
426: A = baij->A; spA = (Mat_SeqBAIJ*)A->data; A_val = spA->a;
427: B = baij->B; spB = (Mat_SeqBAIJ*)B->data; B_val = spB->a;
428: nz = spA->nz + spB->nz; /* total nonzero entries of mat */
429: if (!baij->colmap) {
430: MatCreateColmap_MPIBAIJ_Private(mat);
431: }
432: colmap = baij->colmap;
433: MatGetColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
434: MatGetColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);
436: if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
437: PetscInt *garray;
438: PetscMalloc1(B->cmap->n,&garray);
439: for (i=0; i<baij->B->cmap->n/bs; i++) {
440: for (j=0; j<bs; j++) {
441: garray[i*bs+j] = bs*baij->garray[i]+j;
442: }
443: }
444: VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,garray,&c->vscale);
445: VecBindToCPU(c->vscale,PETSC_TRUE);
446: PetscFree(garray);
447: }
448: } else if (isSELL) {
449: Mat_MPISELL *sell=(Mat_MPISELL*)mat->data;
450: Mat_SeqSELL *spA,*spB;
451: A = sell->A; spA = (Mat_SeqSELL*)A->data; A_val = spA->val;
452: B = sell->B; spB = (Mat_SeqSELL*)B->data; B_val = spB->val;
453: nz = spA->nz + spB->nz; /* total nonzero entries of mat */
454: if (!sell->colmap) {
455: /* Allow access to data structures of local part of matrix
456: - creates aij->colmap which maps global column number to local number in part B */
457: MatCreateColmap_MPISELL_Private(mat);
458: }
459: colmap = sell->colmap;
460: MatGetColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
461: MatGetColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);
463: bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */
465: if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
466: VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,sell->garray,&c->vscale);
467: VecBindToCPU(c->vscale,PETSC_TRUE);
468: }
469: } else {
470: Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
471: Mat_SeqAIJ *spA,*spB;
472: A = aij->A; spA = (Mat_SeqAIJ*)A->data; A_val = spA->a;
473: B = aij->B; spB = (Mat_SeqAIJ*)B->data; B_val = spB->a;
474: nz = spA->nz + spB->nz; /* total nonzero entries of mat */
475: if (!aij->colmap) {
476: /* Allow access to data structures of local part of matrix
477: - creates aij->colmap which maps global column number to local number in part B */
478: MatCreateColmap_MPIAIJ_Private(mat);
479: }
480: colmap = aij->colmap;
481: MatGetColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
482: MatGetColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);
484: bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */
486: if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
487: VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,aij->garray,&c->vscale);
488: VecBindToCPU(c->vscale,PETSC_TRUE);
489: }
490: }
492: m = mat->rmap->n/bs;
493: cstart = mat->cmap->rstart/bs;
494: cend = mat->cmap->rend/bs;
496: PetscMalloc2(nis,&c->ncolumns,nis,&c->columns);
497: PetscMalloc1(nis,&c->nrows);
498: PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));
500: if (c->htype[0] == 'd') {
501: PetscMalloc1(nz,&Jentry);
502: PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));
503: c->matentry = Jentry;
504: } else if (c->htype[0] == 'w') {
505: PetscMalloc1(nz,&Jentry2);
506: PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry2));
507: c->matentry2 = Jentry2;
508: } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"htype is not supported");
510: PetscMalloc2(m+1,&rowhit,m+1,&valaddrhit);
511: nz = 0;
512: ISColoringGetIS(iscoloring,PETSC_OWN_POINTER, PETSC_IGNORE,&c->isa);
514: if (ctype == IS_COLORING_GLOBAL) {
515: MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size);
516: PetscMalloc2(size,&ncolsonproc,size,&disp);
517: }
519: for (i=0; i<nis; i++) { /* for each local color */
520: ISGetLocalSize(c->isa[i],&n);
521: ISGetIndices(c->isa[i],&is);
523: c->ncolumns[i] = n; /* local number of columns of this color on this process */
524: c->columns[i] = (PetscInt*)is;
526: if (ctype == IS_COLORING_GLOBAL) {
527: /* Determine nctot, the total (parallel) number of columns of this color */
528: /* ncolsonproc[j]: local ncolumns on proc[j] of this color */
529: PetscMPIIntCast(n,&nn);
530: MPI_Allgather(&nn,1,MPI_INT,ncolsonproc,1,MPI_INT,PetscObjectComm((PetscObject)mat));
531: nctot = 0; for (j=0; j<size; j++) nctot += ncolsonproc[j];
532: if (!nctot) {
533: PetscInfo(mat,"Coloring of matrix has some unneeded colors with no corresponding rows\n");
534: }
536: disp[0] = 0;
537: for (j=1; j<size; j++) {
538: disp[j] = disp[j-1] + ncolsonproc[j-1];
539: }
541: /* Get cols, the complete list of columns for this color on each process */
542: PetscMalloc1(nctot+1,&cols);
543: MPI_Allgatherv((void*)is,n,MPIU_INT,cols,ncolsonproc,disp,MPIU_INT,PetscObjectComm((PetscObject)mat));
544: } else if (ctype == IS_COLORING_LOCAL) {
545: /* Determine local number of columns of this color on this process, including ghost points */
546: nctot = n;
547: cols = (PetscInt*)is;
548: } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not provided for this MatFDColoring type");
550: /* Mark all rows affect by these columns */
551: PetscArrayzero(rowhit,m);
552: bs2 = bs*bs;
553: nrows_i = 0;
554: for (j=0; j<nctot; j++) { /* loop over columns*/
555: if (ctype == IS_COLORING_LOCAL) {
556: col = ltog[cols[j]];
557: } else {
558: col = cols[j];
559: }
560: if (col >= cstart && col < cend) { /* column is in A, diagonal block of mat */
561: tmp = A_ci[col-cstart];
562: row = A_cj + tmp;
563: nrows = A_ci[col-cstart+1] - tmp;
564: nrows_i += nrows;
566: /* loop over columns of A marking them in rowhit */
567: for (k=0; k<nrows; k++) {
568: /* set valaddrhit for part A */
569: spidx = bs2*spidxA[tmp + k];
570: valaddrhit[*row] = &A_val[spidx];
571: rowhit[*row++] = col - cstart + 1; /* local column index */
572: }
573: } else { /* column is in B, off-diagonal block of mat */
574: #if defined(PETSC_USE_CTABLE)
575: PetscTableFind(colmap,col+1,&colb);
576: colb--;
577: #else
578: colb = colmap[col] - 1; /* local column index */
579: #endif
580: if (colb == -1) {
581: nrows = 0;
582: } else {
583: colb = colb/bs;
584: tmp = B_ci[colb];
585: row = B_cj + tmp;
586: nrows = B_ci[colb+1] - tmp;
587: }
588: nrows_i += nrows;
589: /* loop over columns of B marking them in rowhit */
590: for (k=0; k<nrows; k++) {
591: /* set valaddrhit for part B */
592: spidx = bs2*spidxB[tmp + k];
593: valaddrhit[*row] = &B_val[spidx];
594: rowhit[*row++] = colb + 1 + cend - cstart; /* local column index */
595: }
596: }
597: }
598: c->nrows[i] = nrows_i;
600: if (c->htype[0] == 'd') {
601: for (j=0; j<m; j++) {
602: if (rowhit[j]) {
603: Jentry[nz].row = j; /* local row index */
604: Jentry[nz].col = rowhit[j] - 1; /* local column index */
605: Jentry[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */
606: nz++;
607: }
608: }
609: } else { /* c->htype == 'wp' */
610: for (j=0; j<m; j++) {
611: if (rowhit[j]) {
612: Jentry2[nz].row = j; /* local row index */
613: Jentry2[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */
614: nz++;
615: }
616: }
617: }
618: if (ctype == IS_COLORING_GLOBAL) {
619: PetscFree(cols);
620: }
621: }
622: if (ctype == IS_COLORING_GLOBAL) {
623: PetscFree2(ncolsonproc,disp);
624: }
626: if (bcols > 1) { /* reorder Jentry for faster MatFDColoringApply() */
627: MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz);
628: }
630: if (isBAIJ) {
631: MatRestoreColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
632: MatRestoreColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);
633: PetscMalloc1(bs*mat->rmap->n,&c->dy);
634: } else if (isSELL) {
635: MatRestoreColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
636: MatRestoreColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);
637: } else {
638: MatRestoreColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);
639: MatRestoreColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);
640: }
642: ISColoringRestoreIS(iscoloring,PETSC_OWN_POINTER,&c->isa);
643: PetscFree2(rowhit,valaddrhit);
645: if (ctype == IS_COLORING_LOCAL) {
646: ISLocalToGlobalMappingRestoreIndices(map,<og);
647: }
648: PetscInfo3(c,"ncolors %D, brows %D and bcols %D are used.\n",c->ncolors,c->brows,c->bcols);
649: return(0);
650: }
652: PetscErrorCode MatFDColoringCreate_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
653: {
655: PetscInt bs,nis=iscoloring->n,m=mat->rmap->n;
656: PetscBool isBAIJ,isSELL;
659: /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian;
660: bcols is chosen s.t. dy-array takes 50% of memory space as mat */
661: MatGetBlockSize(mat,&bs);
662: PetscObjectBaseTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);
663: PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL);
664: if (isBAIJ || m == 0) {
665: c->brows = m;
666: c->bcols = 1;
667: } else if (isSELL) {
668: /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
669: Mat_MPISELL *sell=(Mat_MPISELL*)mat->data;
670: Mat_SeqSELL *spA,*spB;
671: Mat A,B;
672: PetscInt nz,brows,bcols;
673: PetscReal mem;
675: bs = 1; /* only bs=1 is supported for MPISELL matrix */
677: A = sell->A; spA = (Mat_SeqSELL*)A->data;
678: B = sell->B; spB = (Mat_SeqSELL*)B->data;
679: nz = spA->nz + spB->nz; /* total local nonzero entries of mat */
680: mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
681: bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
682: brows = 1000/bcols;
683: if (bcols > nis) bcols = nis;
684: if (brows == 0 || brows > m) brows = m;
685: c->brows = brows;
686: c->bcols = bcols;
687: } else { /* mpiaij matrix */
688: /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
689: Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
690: Mat_SeqAIJ *spA,*spB;
691: Mat A,B;
692: PetscInt nz,brows,bcols;
693: PetscReal mem;
695: bs = 1; /* only bs=1 is supported for MPIAIJ matrix */
697: A = aij->A; spA = (Mat_SeqAIJ*)A->data;
698: B = aij->B; spB = (Mat_SeqAIJ*)B->data;
699: nz = spA->nz + spB->nz; /* total local nonzero entries of mat */
700: mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
701: bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
702: brows = 1000/bcols;
703: if (bcols > nis) bcols = nis;
704: if (brows == 0 || brows > m) brows = m;
705: c->brows = brows;
706: c->bcols = bcols;
707: }
709: c->M = mat->rmap->N/bs; /* set the global rows and columns and local rows */
710: c->N = mat->cmap->N/bs;
711: c->m = mat->rmap->n/bs;
712: c->rstart = mat->rmap->rstart/bs;
713: c->ncolors = nis;
714: return(0);
715: }
717: /*@C
719: MatFDColoringSetValues - takes a matrix in compressed color format and enters the matrix into a PETSc Mat
721: Collective on J
723: Input Parameters:
724: + J - the sparse matrix
725: . coloring - created with MatFDColoringCreate() and a local coloring
726: - y - column major storage of matrix values with one color of values per column, the number of rows of y should match
727: the number of local rows of J and the number of columns is the number of colors.
729: Level: intermediate
731: Notes: the matrix in compressed color format may come from an Automatic Differentiation code
733: The code will be slightly faster if MatFDColoringSetBlockSize(coloring,PETSC_DEFAULT,nc); is called immediately after creating the coloring
735: .seealso: MatFDColoringCreate(), ISColoring, ISColoringCreate(), ISColoringSetType(), IS_COLORING_LOCAL, MatFDColoringSetBlockSize()
737: @*/
738: PetscErrorCode MatFDColoringSetValues(Mat J,MatFDColoring coloring,const PetscScalar *y)
739: {
740: PetscErrorCode ierr;
741: MatEntry2 *Jentry2;
742: PetscInt row,i,nrows_k,l,ncolors,nz = 0,bcols,nbcols = 0;
743: const PetscInt *nrows;
744: PetscBool eq;
749: PetscObjectCompareId((PetscObject)J,coloring->matid,&eq);
750: if (!eq) SETERRQ(PetscObjectComm((PetscObject)J),PETSC_ERR_ARG_WRONG,"Matrix used with MatFDColoringSetValues() must be that used with MatFDColoringCreate()");
751: Jentry2 = coloring->matentry2;
752: nrows = coloring->nrows;
753: ncolors = coloring->ncolors;
754: bcols = coloring->bcols;
756: for (i=0; i<ncolors; i+=bcols) {
757: nrows_k = nrows[nbcols++];
758: for (l=0; l<nrows_k; l++) {
759: row = Jentry2[nz].row; /* local row index */
760: *(Jentry2[nz++].valaddr) = y[row];
761: }
762: y += bcols*coloring->m;
763: }
764: MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);
765: MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);
766: return(0);
767: }