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