xref: /petsc/src/mat/impls/aij/mpi/fdmpiaij.c (revision 5fedf56b96a6cbb1434787ca5e034afc7cfcfb65)
1 
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 #undef __FUNCT__
7 #define __FUNCT__ "MatFDColoringApply_BAIJ"
8 PetscErrorCode  MatFDColoringApply_BAIJ(Mat J,MatFDColoring coloring,Vec x1,void *sctx)
9 {
10   PetscErrorCode    (*f)(void*,Vec,Vec,void*) = (PetscErrorCode (*)(void*,Vec,Vec,void*))coloring->f;
11   PetscErrorCode    ierr;
12   PetscInt          k,cstart,cend,l,row,col,nz,spidx,i,j;
13   PetscScalar       dx=0.0,*w3_array,*dy_i,*dy=coloring->dy;
14   PetscScalar       *vscale_array;
15   const PetscScalar *xx;
16   PetscReal         epsilon=coloring->error_rel,umin=coloring->umin,unorm;
17   Vec               w1=coloring->w1,w2=coloring->w2,w3,vscale=coloring->vscale;
18   void              *fctx=coloring->fctx;
19   PetscInt          ctype=coloring->ctype,nxloc,nrows_k;
20   PetscScalar       *valaddr;
21   MatEntry          *Jentry=coloring->matentry;
22   MatEntry2         *Jentry2=coloring->matentry2;
23   const PetscInt    ncolors=coloring->ncolors,*ncolumns=coloring->ncolumns,*nrows=coloring->nrows;
24   PetscInt          bs=J->rmap->bs;
25 
26   PetscFunctionBegin;
27   /* (1) Set w1 = F(x1) */
28   if (!coloring->fset) {
29     ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
30     ierr = (*f)(sctx,x1,w1,fctx);CHKERRQ(ierr);
31     ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
32   } else {
33     coloring->fset = PETSC_FALSE;
34   }
35 
36   /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */
37   ierr = VecGetLocalSize(x1,&nxloc);CHKERRQ(ierr);
38   if (coloring->htype[0] == 'w') {
39     /* vscale = dx is a constant scalar */
40     ierr = VecNorm(x1,NORM_2,&unorm);CHKERRQ(ierr);
41     dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon);
42   } else {
43     ierr = VecGetArrayRead(x1,&xx);CHKERRQ(ierr);
44     ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr);
45     for (col=0; col<nxloc; col++) {
46       dx = xx[col];
47       if (PetscAbsScalar(dx) < umin) {
48         if (PetscRealPart(dx) >= 0.0)      dx = umin;
49         else if (PetscRealPart(dx) < 0.0 ) dx = -umin;
50       }
51       dx               *= epsilon;
52       vscale_array[col] = 1.0/dx;
53     }
54     ierr = VecRestoreArrayRead(x1,&xx);CHKERRQ(ierr);
55     ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr);
56   }
57   if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') {
58     ierr = VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
59     ierr = VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
60   }
61 
62   /* (3) Loop over each color */
63   if (!coloring->w3) {
64     ierr = VecDuplicate(x1,&coloring->w3);CHKERRQ(ierr);
65     ierr = PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);CHKERRQ(ierr);
66   }
67   w3 = coloring->w3;
68 
69   ierr = VecGetOwnershipRange(x1,&cstart,&cend);CHKERRQ(ierr); /* used by ghosted vscale */
70   if (vscale) {
71     ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr);
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     ierr = VecCopy(x1,w3);CHKERRQ(ierr);
82     dy_i = dy;
83     for (i=0; i<bs; i++) {     /* Loop over a block of columns */
84       ierr = VecGetArray(w3,&w3_array);CHKERRQ(ierr);
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       ierr = VecRestoreArray(w3,&w3_array);CHKERRQ(ierr);
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       ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
109       ierr = VecPlaceArray(w2,dy_i);CHKERRQ(ierr); /* place w2 to the array dy_i */
110       ierr = (*f)(sctx,w3,w2,fctx);CHKERRQ(ierr);
111       ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
112       ierr = VecAXPY(w2,-1.0,w1);CHKERRQ(ierr);
113       ierr = VecResetArray(w2);CHKERRQ(ierr);
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   ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
151   ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
152   if (vscale) {
153     ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr);
154   }
155 
156   coloring->currentcolor = -1;
157   PetscFunctionReturn(0);
158 }
159 
160 #undef __FUNCT__
161 #define __FUNCT__ "MatFDColoringApply_AIJ"
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   PetscErrorCode    ierr;
166   PetscInt          k,cstart,cend,l,row,col,nz;
167   PetscScalar       dx=0.0,*y,*w3_array;
168   const PetscScalar *xx;
169   PetscScalar       *vscale_array;
170   PetscReal         epsilon=coloring->error_rel,umin=coloring->umin,unorm;
171   Vec               w1=coloring->w1,w2=coloring->w2,w3,vscale=coloring->vscale;
172   void              *fctx=coloring->fctx;
173   PetscInt          ctype=coloring->ctype,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 
178   PetscFunctionBegin;
179   /* (1) Set w1 = F(x1) */
180   if (!coloring->fset) {
181     ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
182     ierr = (*f)(sctx,x1,w1,fctx);CHKERRQ(ierr);
183     ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
184   } else {
185     coloring->fset = PETSC_FALSE;
186   }
187 
188   /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */
189   if (coloring->htype[0] == 'w') {
190     /* vscale = 1./dx is a constant scalar */
191     ierr = VecNorm(x1,NORM_2,&unorm);CHKERRQ(ierr);
192     dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon);
193   } else {
194     ierr = VecGetLocalSize(x1,&nxloc);CHKERRQ(ierr);
195     ierr = VecGetArrayRead(x1,&xx);CHKERRQ(ierr);
196     ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr);
197     for (col=0; col<nxloc; col++) {
198       dx = xx[col];
199       if (PetscAbsScalar(dx) < umin) {
200         if (PetscRealPart(dx) >= 0.0)      dx = umin;
201         else if (PetscRealPart(dx) < 0.0 ) dx = -umin;
202       }
203       dx               *= epsilon;
204       vscale_array[col] = 1.0/dx;
205     }
206     ierr = VecRestoreArrayRead(x1,&xx);CHKERRQ(ierr);
207     ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr);
208   }
209   if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') {
210     ierr = VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
211     ierr = VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
212   }
213 
214   /* (3) Loop over each color */
215   if (!coloring->w3) {
216     ierr = VecDuplicate(x1,&coloring->w3);CHKERRQ(ierr);
217     ierr = PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);CHKERRQ(ierr);
218   }
219   w3 = coloring->w3;
220 
221   ierr = VecGetOwnershipRange(x1,&cstart,&cend);CHKERRQ(ierr); /* used by ghosted vscale */
222   if (vscale) {
223     ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr);
224   }
225   nz   = 0;
226 
227   if (coloring->bcols > 1) { /* use blocked insertion of Jentry */
228     PetscInt    i,m=J->rmap->n,nbcols,bcols=coloring->bcols;
229     PetscScalar *dy=coloring->dy,*dy_k;
230 
231     nbcols = 0;
232     for (k=0; k<ncolors; k+=bcols) {
233       coloring->currentcolor = k;
234 
235       /*
236        (3-1) Loop over each column associated with color
237        adding the perturbation to the vector w3 = x1 + dx.
238        */
239 
240       dy_k = dy;
241       if (k + bcols > ncolors) bcols = ncolors - k;
242       for (i=0; i<bcols; i++) {
243 
244         ierr = VecCopy(x1,w3);CHKERRQ(ierr);
245         ierr = VecGetArray(w3,&w3_array);CHKERRQ(ierr);
246         if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */
247         if (coloring->htype[0] == 'w') {
248           for (l=0; l<ncolumns[k+i]; l++) {
249             col = coloring->columns[k+i][l]; /* local column (in global index!) of the matrix we are probing for */
250             w3_array[col] += 1.0/dx;
251           }
252         } else { /* htype == 'ds' */
253           vscale_array -= cstart; /* shift pointer so global index can be used */
254           for (l=0; l<ncolumns[k+i]; l++) {
255             col = coloring->columns[k+i][l]; /* local column (in global index!) of the matrix we are probing for */
256             w3_array[col] += 1.0/vscale_array[col];
257           }
258           vscale_array += cstart;
259         }
260         if (ctype == IS_COLORING_GLOBAL) w3_array += cstart;
261         ierr = VecRestoreArray(w3,&w3_array);CHKERRQ(ierr);
262 
263         /*
264          (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations)
265                            w2 = F(x1 + dx) - F(x1)
266          */
267         ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
268         ierr = VecPlaceArray(w2,dy_k);CHKERRQ(ierr); /* place w2 to the array dy_i */
269         ierr = (*f)(sctx,w3,w2,fctx);CHKERRQ(ierr);
270         ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
271         ierr = VecAXPY(w2,-1.0,w1);CHKERRQ(ierr);
272         ierr = VecResetArray(w2);CHKERRQ(ierr);
273         dy_k += m; /* points to dy+i*nxloc */
274       }
275 
276       /*
277        (3-3) Loop over block rows of vector, putting results into Jacobian matrix
278        */
279       nrows_k = nrows[nbcols++];
280       ierr = VecGetArray(w2,&y);CHKERRQ(ierr);
281 
282       if (coloring->htype[0] == 'w') {
283         for (l=0; l<nrows_k; l++) {
284           row                      = Jentry2[nz].row;   /* local row index */
285           *(Jentry2[nz++].valaddr) = dy[row]*dx;
286         }
287       } else { /* htype == 'ds' */
288         for (l=0; l<nrows_k; l++) {
289           row                   = Jentry[nz].row;   /* local row index */
290           *(Jentry[nz].valaddr) = dy[row]*vscale_array[Jentry[nz].col];
291           nz++;
292         }
293       }
294       ierr = VecRestoreArray(w2,&y);CHKERRQ(ierr);
295     }
296   } else { /* bcols == 1 */
297     for (k=0; k<ncolors; k++) {
298       coloring->currentcolor = k;
299 
300       /*
301        (3-1) Loop over each column associated with color
302        adding the perturbation to the vector w3 = x1 + dx.
303        */
304       ierr = VecCopy(x1,w3);CHKERRQ(ierr);
305       ierr = VecGetArray(w3,&w3_array);CHKERRQ(ierr);
306       if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */
307       if (coloring->htype[0] == 'w') {
308         for (l=0; l<ncolumns[k]; l++) {
309           col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
310           w3_array[col] += 1.0/dx;
311         }
312       } else { /* htype == 'ds' */
313         vscale_array -= cstart; /* shift pointer so global index can be used */
314         for (l=0; l<ncolumns[k]; l++) {
315           col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
316           w3_array[col] += 1.0/vscale_array[col];
317         }
318         vscale_array += cstart;
319       }
320       if (ctype == IS_COLORING_GLOBAL) w3_array += cstart;
321       ierr = VecRestoreArray(w3,&w3_array);CHKERRQ(ierr);
322 
323       /*
324        (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations)
325                            w2 = F(x1 + dx) - F(x1)
326        */
327       ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
328       ierr = (*f)(sctx,w3,w2,fctx);CHKERRQ(ierr);
329       ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
330       ierr = VecAXPY(w2,-1.0,w1);CHKERRQ(ierr);
331 
332       /*
333        (3-3) Loop over rows of vector, putting results into Jacobian matrix
334        */
335       nrows_k = nrows[k];
336       ierr = VecGetArray(w2,&y);CHKERRQ(ierr);
337       if (coloring->htype[0] == 'w') {
338         for (l=0; l<nrows_k; l++) {
339           row                      = Jentry2[nz].row;   /* local row index */
340           *(Jentry2[nz++].valaddr) = y[row]*dx;
341         }
342       } else { /* htype == 'ds' */
343         for (l=0; l<nrows_k; l++) {
344           row                   = Jentry[nz].row;   /* local row index */
345           *(Jentry[nz].valaddr) = y[row]*vscale_array[Jentry[nz].col];
346           nz++;
347         }
348       }
349       ierr = VecRestoreArray(w2,&y);CHKERRQ(ierr);
350     }
351   }
352 
353   ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
354   ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
355   if (vscale) {
356     ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr);
357   }
358   coloring->currentcolor = -1;
359   PetscFunctionReturn(0);
360 }
361 
362 #undef __FUNCT__
363 #define __FUNCT__ "MarkRowsForCol_private"
364 PETSC_STATIC_INLINE PetscErrorCode MarkRowsForCol_private(PetscInt nctot,PetscInt *cols,
365                                    Mat mat,ISColoring iscoloring,MatFDColoring c,const PetscInt *ltog,
366                                    const PetscInt *A_ci,const PetscInt *A_cj,PetscScalar *A_val,
367                                    const PetscInt *B_ci,const PetscInt *B_cj,PetscScalar *B_val,PetscInt *spidxA,PetscInt *spidxB,
368                                    PetscInt *rowhit,PetscScalar **valaddrhit,PetscTable colmap,PetscInt *nrows_i_out)
369 {
370   PetscErrorCode         ierr;
371   PetscInt               ctype=c->ctype,m;
372   PetscInt               j,col,nrows,k,spidx,colb,nrows_i;
373   const PetscInt         *row=NULL;
374   PetscInt               cstart,cend,bs;
375 
376   PetscFunctionBegin;
377   //printf(" \n[%d] MarkRows nctot %d\n",rank,nctot);
378   bs        = 1; /* only bs=1 is supported for non MPIBAIJ matrix */
379   m         = mat->rmap->n/bs;
380   cstart    = mat->cmap->rstart/bs;
381   cend      = mat->cmap->rend/bs;
382 
383   ierr    = PetscMemzero(rowhit,m*sizeof(PetscInt));CHKERRQ(ierr);
384   nrows_i = 0;
385   for (j=0; j<nctot; j++) { /* loop over columns*/
386     if (ctype == IS_COLORING_GHOSTED) {
387       col = ltog[cols[j]];
388     } else {
389       col = cols[j];
390     }
391 
392     if (col >= cstart && col < cend) { /* column is in A, diagonal block of mat */
393       row      = A_cj + A_ci[col-cstart];
394       nrows    = A_ci[col-cstart+1] - A_ci[col-cstart];
395       nrows_i += nrows;
396 
397       /* loop over columns of A marking them in rowhit */
398       for (k=0; k<nrows; k++) {
399         /* set valaddrhit for part A */
400         spidx            = bs*bs*spidxA[A_ci[col-cstart] + k];
401         valaddrhit[*row] = &A_val[spidx];
402         rowhit[*row]   = col - cstart + 1; /* local column index */
403         //printf("[%d] MarkRows valaddrhit[%d] %p, rowhit %d\n",rank,*row, valaddrhit[*row],rowhit[*row] );
404         row++;
405       }
406     } else { /* column is in B, off-diagonal block of mat */
407 #if defined(PETSC_USE_CTABLE)
408       ierr = PetscTableFind(colmap,col+1,&colb);CHKERRQ(ierr);
409       colb--;
410 #else
411       colb = colmap[col] - 1; /* local column index */
412 #endif
413       if (colb == -1) {
414         nrows = 0;
415       } else {
416         colb  = colb/bs;
417         row   = B_cj + B_ci[colb];
418         nrows = B_ci[colb+1] - B_ci[colb];
419       }
420       nrows_i += nrows;
421 
422       /* loop over columns of B marking them in rowhit */
423       for (k=0; k<nrows; k++) {
424         /* set valaddrhit for part B */
425         spidx            = bs*bs*spidxB[B_ci[colb] + k];
426         valaddrhit[*row] = &B_val[spidx];
427         rowhit[*row]   = colb + 1 + cend - cstart; /* local column index */
428         //printf("[%d] valaddrhit[%d] %p, rowhit %d\n",rank,*row, valaddrhit[*row],rowhit[*row] );
429         row++;
430       }
431     }
432   }
433   *nrows_i_out = nrows_i;
434   PetscFunctionReturn(0);
435 }
436 
437 #undef __FUNCT__
438 #define __FUNCT__ "MatFDColoringSetUp_MPIXAIJ"
439 PetscErrorCode MatFDColoringSetUp_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
440 {
441   PetscErrorCode         ierr;
442   PetscMPIInt            size,*ncolsonproc,*disp,nn,rank;
443   PetscInt               i,n,nrows,nrows_i,j,k,m,ncols,col,*rowhit,cstart,cend,colb;
444   const PetscInt         *is,*A_ci,*A_cj,*B_ci,*B_cj,*row=NULL,*ltog=NULL;
445   PetscInt               nis=iscoloring->n,nctot,*cols,*cols_new;
446   IS                     *isa;
447   ISLocalToGlobalMapping map=mat->cmap->mapping;
448   PetscInt               ctype=c->ctype,*spidxA,*spidxB,nz,bs,bs2,spidx;
449   Mat                    A,B;
450   PetscScalar            *A_val,*B_val,**valaddrhit;
451   MatEntry               *Jentry;
452   MatEntry2              *Jentry2;
453   PetscBool              isBAIJ;
454   PetscInt               bcols=c->bcols;
455   MPI_Comm               comm;
456   PetscMPIInt            tag,nrecvs,nsends,proc;
457   MPI_Request            *rwaits = NULL,*swaits = NULL;
458   MPI_Status             status;
459 #if defined(PETSC_USE_CTABLE)
460   PetscTable             colmap=NULL;
461 #else
462   PetscInt               *colmap=NULL;     /* local col number of off-diag col */
463 #endif
464 
465   PetscFunctionBegin;
466   ierr = PetscObjectGetComm((PetscObject)mat,&comm);CHKERRQ(ierr);
467   ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
468 
469   if (ctype == IS_COLORING_GHOSTED) {
470     if (!map) SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_INCOMP,"When using ghosted differencing matrix must have local to global mapping provided with MatSetLocalToGlobalMapping");
471     ierr = ISLocalToGlobalMappingGetIndices(map,&ltog);CHKERRQ(ierr);
472   }
473 
474   ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
475   ierr = PetscObjectTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);CHKERRQ(ierr);
476   if (isBAIJ) {
477     Mat_MPIBAIJ *baij=(Mat_MPIBAIJ*)mat->data;
478     Mat_SeqBAIJ *spA,*spB;
479     A = baij->A;  spA = (Mat_SeqBAIJ*)A->data; A_val = spA->a;
480     B = baij->B;  spB = (Mat_SeqBAIJ*)B->data; B_val = spB->a;
481     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
482     if (!baij->colmap) {
483       ierr = MatCreateColmap_MPIBAIJ_Private(mat);CHKERRQ(ierr);
484     }
485     colmap = baij->colmap;
486     ierr = MatGetColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
487     ierr = MatGetColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
488 
489     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') {  /* create vscale for storing dx */
490       PetscInt    *garray;
491       ierr = PetscMalloc1(B->cmap->n,&garray);CHKERRQ(ierr);
492       for (i=0; i<baij->B->cmap->n/bs; i++) {
493         for (j=0; j<bs; j++) {
494           garray[i*bs+j] = bs*baij->garray[i]+j;
495         }
496       }
497       ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,garray,&c->vscale);CHKERRQ(ierr);
498       ierr = PetscFree(garray);CHKERRQ(ierr);
499     }
500   } else {
501     Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
502     Mat_SeqAIJ *spA,*spB;
503     A = aij->A;  spA = (Mat_SeqAIJ*)A->data; A_val = spA->a;
504     B = aij->B;  spB = (Mat_SeqAIJ*)B->data; B_val = spB->a;
505     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
506     if (!aij->colmap) {
507       /* Allow access to data structures of local part of matrix
508        - creates aij->colmap which maps global column number to local number in part B */
509       ierr = MatCreateColmap_MPIAIJ_Private(mat);CHKERRQ(ierr);
510     }
511     colmap = aij->colmap;
512     ierr = MatGetColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
513     ierr = MatGetColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
514 
515     bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */
516 
517     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
518       ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,aij->garray,&c->vscale);CHKERRQ(ierr);
519     }
520   }
521 
522   m         = mat->rmap->n/bs;
523   cstart    = mat->cmap->rstart/bs;
524   cend      = mat->cmap->rend/bs;
525 
526   ierr       = PetscMalloc1(nis,&c->ncolumns);CHKERRQ(ierr);
527   ierr       = PetscMalloc1(nis,&c->columns);CHKERRQ(ierr);
528   ierr       = PetscMalloc1(nis,&c->nrows);CHKERRQ(ierr);
529   ierr       = PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));CHKERRQ(ierr);
530 
531   if (c->htype[0] == 'd') {
532     ierr       = PetscMalloc1(nz,&Jentry);CHKERRQ(ierr);
533     ierr       = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));CHKERRQ(ierr);
534     c->matentry = Jentry;
535   } else if (c->htype[0] == 'w') {
536     ierr       = PetscMalloc1(nz,&Jentry2);CHKERRQ(ierr);
537     ierr       = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry2));CHKERRQ(ierr);
538     c->matentry2 = Jentry2;
539   } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"htype is not supported");
540 
541   ierr = PetscMalloc2(m+1,&rowhit,m+1,&valaddrhit);CHKERRQ(ierr);
542   nz = 0;
543   ierr = ISColoringGetIS(iscoloring,PETSC_IGNORE,&isa);CHKERRQ(ierr);
544   for (i=0; i<nis; i++) { /* for each local color */
545     ierr = ISGetLocalSize(isa[i],&n);CHKERRQ(ierr);
546     ierr = ISGetIndices(isa[i],&is);CHKERRQ(ierr);
547 
548     c->ncolumns[i] = n; /* local number of columns of this color on this process */
549     if (n) {
550       ierr = PetscMalloc1(n,&c->columns[i]);CHKERRQ(ierr);
551       ierr = PetscLogObjectMemory((PetscObject)c,n*sizeof(PetscInt));CHKERRQ(ierr);
552       ierr = PetscMemcpy(c->columns[i],is,n*sizeof(PetscInt));CHKERRQ(ierr);
553     } else {
554       c->columns[i] = 0;
555     }
556 
557     if (ctype == IS_COLORING_GLOBAL) {
558       /* Determine nctot, the total (parallel) number of columns of this color */
559       ierr = MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size);CHKERRQ(ierr);
560       ierr = PetscMalloc2(size,&ncolsonproc,size,&disp);CHKERRQ(ierr);
561 
562       /* ncolsonproc[j]: local ncolumns on proc[j] of this color */
563       ierr  = PetscMPIIntCast(n,&nn);CHKERRQ(ierr);
564       ierr  = MPI_Allgather(&nn,1,MPI_INT,ncolsonproc,1,MPI_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr);
565       nctot = 0;
566       for (j=0; j<size; j++) nctot += ncolsonproc[j];
567       if (!nctot) {
568         ierr = PetscInfo(mat,"Coloring of matrix has some unneeded colors with no corresponding rows\n");CHKERRQ(ierr);
569       }
570 
571       disp[0] = 0;
572       for (j=1; j<size; j++) {
573         disp[j] = disp[j-1] + ncolsonproc[j-1];
574       }
575 
576       /* Get cols, the complete list of columns for this color on each process */
577       ierr = PetscMalloc1(nctot+1,&cols);CHKERRQ(ierr);
578 
579       /****************** non-scalable !!! *********************/
580       /* ierr = MPI_Allgatherv((void*)is,n,MPIU_INT,cols,ncolsonproc,disp,MPIU_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr); */
581       cols_new = cols;
582 
583       nrecvs = size-1;
584       ierr = PetscCommGetNewTag(comm,&tag);CHKERRQ(ierr);
585       ierr = PetscMalloc2(nrecvs,&rwaits,nrecvs,&swaits);CHKERRQ(ierr);
586 
587       nrecvs = 0;
588       for (proc=0; proc<size; proc++) {
589         if (proc == rank || ncolsonproc[proc] == 0 ) continue;
590         /* printf("[%d] Irecv %d from [%d] \n",rank,ncolsonproc[proc],proc); */
591         ierr = MPI_Irecv(cols_new+disp[proc],ncolsonproc[proc],MPIU_INT,proc,tag,comm,rwaits+nrecvs);CHKERRQ(ierr);
592         nrecvs++;
593       }
594 
595       nsends = 0;
596       for (proc=0; proc<size; proc++) {
597         if (proc == rank || n == 0) continue;
598         /* printf("[%d] Isend %d is to [%d] \n",rank,n,proc); */
599         ierr = MPI_Isend((void*)is,n,MPIU_INT,proc,tag,comm,swaits+nsends);CHKERRQ(ierr);
600         nsends++;
601       }
602 
603       for (j=0; j<ncolsonproc[rank]; j++) {
604         cols_new[disp[rank] + j] = is[j];
605       }
606 
607       j = nrecvs;
608       while (j--) {
609         ierr = MPI_Waitany(nrecvs,rwaits,&k,&status);CHKERRQ(ierr);
610       }
611       if (nsends) {ierr = MPI_Waitall(nrecvs,swaits,MPI_STATUSES_IGNORE);CHKERRQ(ierr);}
612       ierr = PetscFree2(rwaits,swaits);CHKERRQ(ierr);
613 
614       /******************/
615       ierr = PetscFree2(ncolsonproc,disp);CHKERRQ(ierr);
616     } else if (ctype == IS_COLORING_GHOSTED) {
617       /* Determine local number of columns of this color on this process, including ghost points */
618       nctot = n;
619       ierr  = PetscMalloc1(nctot+1,&cols);CHKERRQ(ierr);
620       ierr  = PetscMemcpy(cols,is,n*sizeof(PetscInt));CHKERRQ(ierr);
621     } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not provided for this MatFDColoring type");
622 
623     /* Mark all rows affect by these columns */
624     /********** new ********/
625     ierr = MarkRowsForCol_private(nctot,cols,mat,iscoloring,c,ltog,A_ci,A_cj,A_val,B_ci,B_cj,B_val,spidxA,spidxB,rowhit,valaddrhit,colmap,&nrows_i);CHKERRQ(ierr);
626 #if 0
627     ierr    = PetscMemzero(rowhit,m*sizeof(PetscInt));CHKERRQ(ierr);
628     bs2     = bs*bs;
629     nrows_i = 0;
630 
631     //printf("  \n[%d] nctot %d\n",rank,nctot);
632     for (j=0; j<nctot; j++) { /* loop over columns*/
633 
634       if (ctype == IS_COLORING_GHOSTED) {
635         col = ltog[cols[j]];
636       } else {
637         col = cols[j];
638       }
639       if (col >= cstart && col < cend) { /* column is in A, diagonal block of mat */
640         row      = A_cj + A_ci[col-cstart];
641         nrows    = A_ci[col-cstart+1] - A_ci[col-cstart];
642         nrows_i += nrows;
643 
644         /* loop over columns of A marking them in rowhit */
645         for (k=0; k<nrows; k++) {
646           /* set valaddrhit for part A */
647           spidx            = bs2*spidxA[A_ci[col-cstart] + k];
648           valaddrhit[*row] = &A_val[spidx];
649           rowhit[*row]   = col - cstart + 1; /* local column index */
650           //printf("[%d] valaddrhit[%d] %p, rowhit %d\n",rank,*row, valaddrhit[*row],rowhit[*row] );
651           row++;
652         }
653       } else { /* column is in B, off-diagonal block of mat */
654 #if defined(PETSC_USE_CTABLE)
655         ierr = PetscTableFind(colmap,col+1,&colb);CHKERRQ(ierr);
656         colb--;
657 #else
658         colb = colmap[col] - 1; /* local column index */
659 #endif
660         if (colb == -1) {
661           nrows = 0;
662         } else {
663           colb  = colb/bs;
664           row   = B_cj + B_ci[colb];
665           nrows = B_ci[colb+1] - B_ci[colb];
666         }
667         nrows_i += nrows;
668 
669         /* loop over columns of B marking them in rowhit */
670         for (k=0; k<nrows; k++) {
671           /* set valaddrhit for part B */
672           spidx            = bs2*spidxB[B_ci[colb] + k];
673           valaddrhit[*row] = &B_val[spidx];
674           rowhit[*row]   = colb + 1 + cend - cstart; /* local column index */
675           //printf("[%d] valaddrhit[%d] %p, rowhit %d\n",rank,*row, valaddrhit[*row],rowhit[*row] );
676           row++;
677         }
678       }
679     } //endif loop
680     /********************************** */
681 #endif
682     ierr = PetscFree(cols);CHKERRQ(ierr);
683     c->nrows[i] = nrows_i;
684 
685     if (c->htype[0] == 'd') {
686       for (j=0; j<m; j++) {
687         if (rowhit[j]) {
688           Jentry[nz].row     = j;              /* local row index */
689           Jentry[nz].col     = rowhit[j] - 1;  /* local column index */
690           Jentry[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
691           nz++;
692         }
693       }
694     } else { /* c->htype == 'wp' */
695       for (j=0; j<m; j++) {
696         if (rowhit[j]) {
697           Jentry2[nz].row     = j;              /* local row index */
698           Jentry2[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
699           nz++;
700         }
701       }
702     }
703   }
704 
705   if (bcols > 1) { /* reorder Jentry for faster MatFDColoringApply() */
706     ierr = MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz);CHKERRQ(ierr);
707   }
708 
709   if (isBAIJ) {
710     ierr = MatRestoreColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
711     ierr = MatRestoreColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
712     ierr = PetscMalloc1(bs*mat->rmap->n,&c->dy);CHKERRQ(ierr);
713   } else {
714     ierr = MatRestoreColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
715     ierr = MatRestoreColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
716   }
717 
718   ierr = ISColoringRestoreIS(iscoloring,&isa);CHKERRQ(ierr);
719   ierr = PetscFree2(rowhit,valaddrhit);CHKERRQ(ierr);
720 
721   if (ctype == IS_COLORING_GHOSTED) {
722     ierr = ISLocalToGlobalMappingRestoreIndices(map,&ltog);CHKERRQ(ierr);
723   }
724   ierr = PetscInfo3(c,"ncolors %D, brows %D and bcols %D are used.\n",c->ncolors,c->brows,c->bcols);CHKERRQ(ierr);
725   PetscFunctionReturn(0);
726 }
727 
728 #undef __FUNCT__
729 #define __FUNCT__ "MatFDColoringCreate_MPIXAIJ"
730 PetscErrorCode MatFDColoringCreate_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
731 {
732   PetscErrorCode ierr;
733   PetscInt       bs,nis=iscoloring->n,m=mat->rmap->n;
734   PetscBool      isBAIJ;
735 
736   PetscFunctionBegin;
737   /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian;
738    bcols is chosen s.t. dy-array takes 50% of memory space as mat */
739   ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
740   ierr = PetscObjectTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);CHKERRQ(ierr);
741   if (isBAIJ || m == 0) {
742     c->brows = m;
743     c->bcols = 1;
744   } else { /* mpiaij matrix */
745     /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
746     Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
747     Mat_SeqAIJ *spA,*spB;
748     Mat        A,B;
749     PetscInt   nz,brows,bcols;
750     PetscReal  mem;
751 
752     bs    = 1; /* only bs=1 is supported for MPIAIJ matrix */
753 
754     A = aij->A;  spA = (Mat_SeqAIJ*)A->data;
755     B = aij->B;  spB = (Mat_SeqAIJ*)B->data;
756     nz = spA->nz + spB->nz; /* total local nonzero entries of mat */
757     mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
758     bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
759     brows = 1000/bcols;
760     if (bcols > nis) bcols = nis;
761     if (brows == 0 || brows > m) brows = m;
762     c->brows = brows;
763     c->bcols = bcols;
764   }
765 
766   c->M       = mat->rmap->N/bs;         /* set the global rows and columns and local rows */
767   c->N       = mat->cmap->N/bs;
768   c->m       = mat->rmap->n/bs;
769   c->rstart  = mat->rmap->rstart/bs;
770   c->ncolors = nis;
771   PetscFunctionReturn(0);
772 }
773