xref: /petsc/src/mat/impls/aij/mpi/fdmpiaij.c (revision 684f2004e99e45d6318bc6d6b71aeefd2c18360e)
1 
2 #include <../src/mat/impls/aij/mpi/mpiaij.h>
3 #include <../src/mat/impls/baij/mpi/mpibaij.h>
4 
5 #undef __FUNCT__
6 #define __FUNCT__ "MatFDColoringApply_BAIJ"
7 PetscErrorCode  MatFDColoringApply_BAIJ(Mat J,MatFDColoring coloring,Vec x1,MatStructure *flag,void *sctx)
8 {
9   PetscErrorCode (*f)(void*,Vec,Vec,void*) = (PetscErrorCode (*)(void*,Vec,Vec,void*))coloring->f;
10   PetscErrorCode ierr;
11   PetscInt       k,cstart,cend,l,row,col,nz,spidx,i,j;
12   PetscScalar    dx=0.0,*y,*xx,*w3_array,*dy_i,*dy=coloring->dy;
13   PetscScalar    *vscale_array;
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   const PetscInt ncolors=coloring->ncolors,*ncolumns=coloring->ncolumns,*nrows=coloring->nrows;
21   PetscInt       bs=J->rmap->bs;
22 
23   PetscFunctionBegin;
24   /* create vscale for storing dx */
25   if (!vscale) {
26     if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') {
27       /* contain the "diagonal" on processor scalings followed by the off processor* - garray must be non-bloced! */
28       Mat_MPIBAIJ    *baij=(Mat_MPIBAIJ*)J->data;
29       PetscInt       *garray;
30       ierr = PetscMalloc(baij->B->cmap->n*sizeof(PetscInt),&garray);CHKERRQ(ierr);
31       for (i=0; i<baij->B->cmap->n/bs; i++) {
32         for (j=0; j<bs; j++) {
33           garray[i*bs+j] = bs*baij->garray[i]+j;
34         }
35       }
36       ierr = VecCreateGhost(PetscObjectComm((PetscObject)J),J->cmap->n,PETSC_DETERMINE,baij->B->cmap->n,garray,&vscale);CHKERRQ(ierr);
37       ierr = PetscFree(garray);CHKERRQ(ierr);
38     } else if (ctype == IS_COLORING_GHOSTED) {
39       ierr = VecDuplicate(x1,&vscale);CHKERRQ(ierr);
40     }
41     coloring->vscale = vscale;
42   }
43 
44   /* (1) Set w1 = F(x1) */
45   if (!coloring->fset) {
46     ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
47     ierr = (*f)(sctx,x1,w1,fctx);CHKERRQ(ierr);
48     ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
49   } else {
50     coloring->fset = PETSC_FALSE;
51   }
52 
53   /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */
54   ierr = VecGetLocalSize(x1,&nxloc);CHKERRQ(ierr);
55   if (coloring->htype[0] == 'w') {
56     /* vscale = dx is a constant scalar */
57     ierr = VecNorm(x1,NORM_2,&unorm);CHKERRQ(ierr);
58     dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon);
59   } else {
60     ierr = VecGetArray(x1,&xx);CHKERRQ(ierr);
61     ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr);
62     for (col=0; col<nxloc; col++) {
63       dx = xx[col];
64       if (PetscAbsScalar(dx) < umin) {
65         if (PetscRealPart(dx) >= 0.0)      dx = umin;
66         else if (PetscRealPart(dx) < 0.0 ) dx = -umin;
67       }
68       dx               *= epsilon;
69       vscale_array[col] = 1.0/dx;
70     }
71     ierr = VecRestoreArray(x1,&xx);CHKERRQ(ierr);
72     ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr);
73   }
74   if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') {
75     ierr = VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
76     ierr = VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
77   }
78 
79   /* (3) Loop over each color */
80   if (!coloring->w3) {
81     ierr = VecDuplicate(x1,&coloring->w3);CHKERRQ(ierr);
82     ierr = PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);CHKERRQ(ierr);
83   }
84   w3 = coloring->w3;
85 
86   ierr = VecGetOwnershipRange(x1,&cstart,&cend);CHKERRQ(ierr); /* used by ghosted vscale */
87   if (vscale) {
88     ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr);
89   }
90   nz   = 0;
91   for (k=0; k<ncolors; k++) {
92     coloring->currentcolor = k;
93 
94     /*
95       (3-1) Loop over each column associated with color
96       adding the perturbation to the vector w3 = x1 + dx.
97     */
98     ierr = VecCopy(x1,w3);CHKERRQ(ierr);
99     dy_i = dy;
100     for (i=0; i<bs; i++) {     /* Loop over a block of columns */
101       ierr = VecGetArray(w3,&w3_array);CHKERRQ(ierr);
102       if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */
103       if (coloring->htype[0] == 'w') {
104         for (l=0; l<ncolumns[k]; l++) {
105           col            = i + bs*coloring->columns[k][l];  /* local column (in global index!) of the matrix we are probing for */
106           w3_array[col] += 1.0/dx;
107           if (i) w3_array[col-1] -= 1.0/dx; /* resume original w3[col-1] */
108         }
109       } else { /* htype == 'ds' */
110         vscale_array -= cstart; /* shift pointer so global index can be used */
111         for (l=0; l<ncolumns[k]; l++) {
112           col = i + bs*coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
113           w3_array[col] += 1.0/vscale_array[col];
114           if (i) w3_array[col-1] -=  1.0/vscale_array[col-1]; /* resume original w3[col-1] */
115         }
116         vscale_array += cstart;
117       }
118       if (ctype == IS_COLORING_GLOBAL) w3_array += cstart;
119       ierr = VecRestoreArray(w3,&w3_array);CHKERRQ(ierr);
120 
121       /*
122        (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations)
123                            w2 = F(x1 + dx) - F(x1)
124        */
125       ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
126       ierr = VecPlaceArray(w2,dy_i);CHKERRQ(ierr); /* place w2 to the array dy_i */
127       ierr = (*f)(sctx,w3,w2,fctx);CHKERRQ(ierr);
128       ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
129       ierr = VecAXPY(w2,-1.0,w1);CHKERRQ(ierr);
130       ierr = VecResetArray(w2);CHKERRQ(ierr);
131       dy_i += nxloc; /* points to dy+i*nxloc */
132     }
133 
134     /*
135      (3-3) Loop over rows of vector, putting results into Jacobian matrix
136     */
137     nrows_k = nrows[k];
138     ierr = VecGetArray(w2,&y);CHKERRQ(ierr);
139     if (coloring->htype[0] == 'w') {
140       for (l=0; l<nrows_k; l++) {
141         row     = bs*Jentry[nz].row;   /* local row index */
142         valaddr = Jentry[nz++].valaddr;
143         spidx   = 0;
144         dy_i    = dy;
145         for (i=0; i<bs; i++) {   /* column of the block */
146           for (j=0; j<bs; j++) { /* row of the block */
147             valaddr[spidx++] = dy_i[row+j]*dx;
148           }
149           dy_i += nxloc; /* points to dy+i*nxloc */
150         }
151       }
152     } else { /* htype == 'ds' */
153       for (l=0; l<nrows_k; l++) {
154         row     = bs*Jentry[nz].row;   /* local row index */
155         col     = bs*Jentry[nz].col;   /* local column index */
156         valaddr = Jentry[nz++].valaddr;
157         spidx   = 0;
158         dy_i    = dy;
159         for (i=0; i<bs; i++) {   /* column of the block */
160           for (j=0; j<bs; j++) { /* row of the block */
161             valaddr[spidx++] = dy_i[row+j]*vscale_array[col+i];
162           }
163           dy_i += nxloc; /* points to dy+i*nxloc */
164         }
165       }
166     }
167     ierr = VecRestoreArray(w2,&y);CHKERRQ(ierr);
168   }
169   ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
170   ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
171   if (vscale) {
172     ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr);
173   }
174 
175   coloring->currentcolor = -1;
176   PetscFunctionReturn(0);
177 }
178 
179 #undef __FUNCT__
180 #define __FUNCT__ "MatFDColoringApply_AIJ"
181 PetscErrorCode  MatFDColoringApply_AIJ(Mat J,MatFDColoring coloring,Vec x1,MatStructure *flag,void *sctx)
182 {
183   PetscErrorCode (*f)(void*,Vec,Vec,void*) = (PetscErrorCode (*)(void*,Vec,Vec,void*))coloring->f;
184   PetscErrorCode ierr;
185   PetscInt       k,cstart,cend,l,row,col,nz;
186   PetscScalar    dx=0.0,*y,*xx,*w3_array;
187   PetscScalar    *vscale_array;
188   PetscReal      epsilon=coloring->error_rel,umin=coloring->umin,unorm;
189   Vec            w1=coloring->w1,w2=coloring->w2,w3,vscale=coloring->vscale;
190   void           *fctx=coloring->fctx;
191   PetscInt       ctype=coloring->ctype,nxloc,nrows_k;
192   MatEntry       *Jentry=coloring->matentry;
193   const PetscInt ncolors=coloring->ncolors,*ncolumns=coloring->ncolumns,*nrows=coloring->nrows;
194 
195   PetscFunctionBegin;
196   /* create vscale for storing dx */
197   if (!vscale) {
198     if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') {
199       Mat_MPIAIJ     *aij=(Mat_MPIAIJ*)J->data;
200       ierr = VecCreateGhost(PetscObjectComm((PetscObject)J),J->cmap->n,PETSC_DETERMINE,aij->B->cmap->n,aij->garray,&vscale);CHKERRQ(ierr);
201     } else if (ctype == IS_COLORING_GHOSTED) {
202       ierr = VecDuplicate(x1,&vscale);CHKERRQ(ierr);
203     }
204     coloring->vscale = vscale;
205   }
206 
207   /* (1) Set w1 = F(x1) */
208   if (!coloring->fset) {
209     ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
210     ierr = (*f)(sctx,x1,w1,fctx);CHKERRQ(ierr);
211     ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
212   } else {
213     coloring->fset = PETSC_FALSE;
214   }
215 
216   /* (2) Compute vscale = 1./dx - the local scale factors, including ghost points */
217   if (coloring->htype[0] == 'w') {
218     /* vscale = 1./dx is a constant scalar */
219     ierr = VecNorm(x1,NORM_2,&unorm);CHKERRQ(ierr);
220     dx = 1.0/(PetscSqrtReal(1.0 + unorm)*epsilon);
221   } else {
222     ierr = VecGetLocalSize(x1,&nxloc);CHKERRQ(ierr);
223     ierr = VecGetArray(x1,&xx);CHKERRQ(ierr);
224     ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr);
225     for (col=0; col<nxloc; col++) {
226       dx = xx[col];
227       if (PetscAbsScalar(dx) < umin) {
228         if (PetscRealPart(dx) >= 0.0)      dx = umin;
229         else if (PetscRealPart(dx) < 0.0 ) dx = -umin;
230       }
231       dx               *= epsilon;
232       vscale_array[col] = 1.0/dx;
233     }
234     ierr = VecRestoreArray(x1,&xx);CHKERRQ(ierr);
235     ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr);
236   }
237   if (ctype == IS_COLORING_GLOBAL && coloring->htype[0] == 'd') {
238     ierr = VecGhostUpdateBegin(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
239     ierr = VecGhostUpdateEnd(vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
240   }
241 
242   /* (3) Loop over each color */
243   if (!coloring->w3) {
244     ierr = VecDuplicate(x1,&coloring->w3);CHKERRQ(ierr);
245     ierr = PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);CHKERRQ(ierr);
246   }
247   w3 = coloring->w3;
248 
249   ierr = VecGetOwnershipRange(x1,&cstart,&cend);CHKERRQ(ierr); /* used by ghosted vscale */
250   if (vscale) {
251     ierr = VecGetArray(vscale,&vscale_array);CHKERRQ(ierr);
252   }
253   nz   = 0;
254   for (k=0; k<ncolors; k++) {
255     coloring->currentcolor = k;
256 
257     /*
258       (3-1) Loop over each column associated with color
259       adding the perturbation to the vector w3 = x1 + dx.
260     */
261     ierr = VecCopy(x1,w3);CHKERRQ(ierr);
262     ierr = VecGetArray(w3,&w3_array);CHKERRQ(ierr);
263     if (ctype == IS_COLORING_GLOBAL) w3_array -= cstart; /* shift pointer so global index can be used */
264     if (coloring->htype[0] == 'w') {
265       for (l=0; l<ncolumns[k]; l++) {
266         col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
267         w3_array[col] += 1.0/dx;
268       }
269     } else { /* htype == 'ds' */
270       vscale_array -= cstart; /* shift pointer so global index can be used */
271       for (l=0; l<ncolumns[k]; l++) {
272         col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
273         w3_array[col] += 1.0/vscale_array[col];
274       }
275       vscale_array += cstart;
276     }
277     if (ctype == IS_COLORING_GLOBAL) w3_array += cstart;
278     ierr = VecRestoreArray(w3,&w3_array);CHKERRQ(ierr);
279 
280     /*
281       (3-2) Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations)
282                            w2 = F(x1 + dx) - F(x1)
283     */
284     ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
285     ierr = (*f)(sctx,w3,w2,fctx);CHKERRQ(ierr);
286     ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
287     ierr = VecAXPY(w2,-1.0,w1);CHKERRQ(ierr);
288 
289     /*
290      (3-3) Loop over rows of vector, putting results into Jacobian matrix
291     */
292     nrows_k = nrows[k];
293     ierr = VecGetArray(w2,&y);CHKERRQ(ierr);
294     if (coloring->htype[0] == 'w') {
295       for (l=0; l<nrows_k; l++) {
296         row                     = Jentry[nz].row;   /* local row index */
297         *(Jentry[nz++].valaddr) = y[row]*dx;
298       }
299     } else { /* htype == 'ds' */
300       for (l=0; l<nrows_k; l++) {
301         row                   = Jentry[nz].row;   /* local row index */
302         *(Jentry[nz].valaddr) = y[row]*vscale_array[Jentry[nz].col];
303         nz++;
304       }
305     }
306     ierr = VecRestoreArray(w2,&y);CHKERRQ(ierr);
307   }
308   ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
309   ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
310   if (vscale) {
311     ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr);
312   }
313 
314   coloring->currentcolor = -1;
315   PetscFunctionReturn(0);
316 }
317 
318 #undef __FUNCT__
319 #define __FUNCT__ "MatFDColoringCreate_MPIXAIJ"
320 PetscErrorCode MatFDColoringCreate_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
321 {
322   PetscErrorCode         ierr;
323   PetscMPIInt            size,*ncolsonproc,*disp,nn;
324   PetscInt               i,n,nrows,nrows_i,j,k,m,ncols,col,*rowhit,cstart,cend,colb;
325   const PetscInt         *is,*A_ci,*A_cj,*B_ci,*B_cj,*row=NULL,*ltog=NULL;
326   PetscInt               nis=iscoloring->n,nctot,*cols;
327   IS                     *isa;
328   ISLocalToGlobalMapping map=mat->cmap->mapping;
329   PetscInt               ctype=c->ctype,*spidxA,*spidxB,nz,bs,bs2,spidx;
330   Mat                    A,B;
331   PetscScalar            *A_val,*B_val,**valaddrhit;
332   MatEntry               *Jentry;
333   PetscBool              isBAIJ;
334 #if defined(PETSC_USE_CTABLE)
335   PetscTable             colmap=NULL;
336 #else
337   PetscInt               *colmap=NULL;     /* local col number of off-diag col */
338 #endif
339 
340   PetscFunctionBegin;
341   if (ctype == IS_COLORING_GHOSTED) {
342     if (!map) SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_INCOMP,"When using ghosted differencing matrix must have local to global mapping provided with MatSetLocalToGlobalMapping");
343     ierr = ISLocalToGlobalMappingGetIndices(map,&ltog);CHKERRQ(ierr);
344   }
345 
346   ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
347   ierr = PetscObjectTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);CHKERRQ(ierr);
348   if (isBAIJ) {
349     Mat_MPIBAIJ *aij=(Mat_MPIBAIJ*)mat->data;
350     Mat_SeqBAIJ *spA,*spB;
351     A = aij->A;  spA = (Mat_SeqBAIJ*)A->data; A_val = spA->a;
352     B = aij->B;  spB = (Mat_SeqBAIJ*)B->data; B_val = spB->a;
353     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
354     if (!aij->colmap) {
355       ierr = MatCreateColmap_MPIBAIJ_Private(mat);CHKERRQ(ierr);
356       colmap = aij->colmap;
357     }
358     ierr = MatGetColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
359     ierr = MatGetColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
360   } else {
361     Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
362     Mat_SeqAIJ *spA,*spB;
363     A = aij->A;  spA = (Mat_SeqAIJ*)A->data; A_val = spA->a;
364     B = aij->B;  spB = (Mat_SeqAIJ*)B->data; B_val = spB->a;
365     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
366     if (!aij->colmap) {
367       /* Allow access to data structures of local part of matrix
368        - creates aij->colmap which maps global column number to local number in part B */
369       ierr = MatCreateColmap_MPIAIJ_Private(mat);CHKERRQ(ierr);
370       colmap = aij->colmap;
371     }
372     ierr = MatGetColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
373     ierr = MatGetColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
374 
375     bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */
376   }
377 
378   m         = mat->rmap->n/bs;
379   cstart    = mat->cmap->rstart/bs;
380   cend      = mat->cmap->rend/bs;
381   c->M      = mat->rmap->N/bs;         /* set the global rows and columns and local rows */
382   c->N      = mat->cmap->N/bs;
383   c->m      = m;
384   c->rstart = mat->rmap->rstart/bs;
385 
386   c->ncolors = nis;
387   ierr       = PetscMalloc(nis*sizeof(PetscInt),&c->ncolumns);CHKERRQ(ierr);
388   ierr       = PetscMalloc(nis*sizeof(PetscInt*),&c->columns);CHKERRQ(ierr);
389   ierr       = PetscMalloc(nis*sizeof(PetscInt),&c->nrows);CHKERRQ(ierr);
390   ierr       = PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));CHKERRQ(ierr);
391 
392   ierr       = PetscMalloc(nz*sizeof(MatEntry),&Jentry);CHKERRQ(ierr);
393   ierr       = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));CHKERRQ(ierr);
394   c->matentry = Jentry;
395 
396   ierr = PetscMalloc2(m+1,PetscInt,&rowhit,m+1,PetscScalar*,&valaddrhit);CHKERRQ(ierr);
397   nz = 0;
398   ierr = ISColoringGetIS(iscoloring,PETSC_IGNORE,&isa);CHKERRQ(ierr);
399   for (i=0; i<nis; i++) { /* for each local color */
400     ierr = ISGetLocalSize(isa[i],&n);CHKERRQ(ierr);
401     ierr = ISGetIndices(isa[i],&is);CHKERRQ(ierr);
402 
403     c->ncolumns[i] = n; /* local number of columns of this color on this process */
404     if (n) {
405       ierr = PetscMalloc(n*sizeof(PetscInt),&c->columns[i]);CHKERRQ(ierr);
406       ierr = PetscLogObjectMemory((PetscObject)c,n*sizeof(PetscInt));CHKERRQ(ierr);
407       ierr = PetscMemcpy(c->columns[i],is,n*sizeof(PetscInt));CHKERRQ(ierr);
408     } else {
409       c->columns[i] = 0;
410     }
411 
412     if (ctype == IS_COLORING_GLOBAL) {
413       /* Determine nctot, the total (parallel) number of columns of this color */
414       ierr = MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size);CHKERRQ(ierr);
415       ierr = PetscMalloc2(size,PetscMPIInt,&ncolsonproc,size,PetscMPIInt,&disp);CHKERRQ(ierr);
416 
417       /* ncolsonproc[j]: local ncolumns on proc[j] of this color */
418       ierr  = PetscMPIIntCast(n,&nn);CHKERRQ(ierr);
419       ierr  = MPI_Allgather(&nn,1,MPI_INT,ncolsonproc,1,MPI_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr);
420       nctot = 0; for (j=0; j<size; j++) nctot += ncolsonproc[j];
421       if (!nctot) {
422         ierr = PetscInfo(mat,"Coloring of matrix has some unneeded colors with no corresponding rows\n");CHKERRQ(ierr);
423       }
424 
425       disp[0] = 0;
426       for (j=1; j<size; j++) {
427         disp[j] = disp[j-1] + ncolsonproc[j-1];
428       }
429 
430       /* Get cols, the complete list of columns for this color on each process */
431       ierr = PetscMalloc((nctot+1)*sizeof(PetscInt),&cols);CHKERRQ(ierr);
432       ierr = MPI_Allgatherv((void*)is,n,MPIU_INT,cols,ncolsonproc,disp,MPIU_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr);
433       ierr = PetscFree2(ncolsonproc,disp);CHKERRQ(ierr);
434     } else if (ctype == IS_COLORING_GHOSTED) {
435       /* Determine local number of columns of this color on this process, including ghost points */
436       nctot = n;
437       ierr  = PetscMalloc((nctot+1)*sizeof(PetscInt),&cols);CHKERRQ(ierr);
438       ierr  = PetscMemcpy(cols,is,n*sizeof(PetscInt));CHKERRQ(ierr);
439     } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not provided for this MatFDColoring type");
440 
441     /* Mark all rows affect by these columns */
442     ierr    = PetscMemzero(rowhit,m*sizeof(PetscInt));CHKERRQ(ierr);
443     bs2     = bs*bs;
444     nrows_i = 0;
445     for (j=0; j<nctot; j++) { /* loop over columns*/
446       if (ctype == IS_COLORING_GHOSTED) {
447         col = ltog[cols[j]];
448       } else {
449         col = cols[j];
450       }
451       if (col >= cstart && col < cend) { /* column is in A, diagonal block of mat */
452         row      = A_cj + A_ci[col-cstart];
453         nrows    = A_ci[col-cstart+1] - A_ci[col-cstart];
454         nrows_i += nrows;
455         /* loop over columns of A marking them in rowhit */
456         for (k=0; k<nrows; k++) {
457           /* set valaddrhit for part A */
458           spidx            = bs2*spidxA[A_ci[col-cstart] + k];
459           valaddrhit[*row] = &A_val[spidx];
460           rowhit[*row++]   = col - cstart + 1; /* local column index */
461         }
462       } else { /* column is in B, off-diagonal block of mat */
463 #if defined(PETSC_USE_CTABLE)
464         ierr = PetscTableFind(colmap,col+1,&colb);CHKERRQ(ierr);
465         colb--;
466 #else
467         colb = colmap[col] - 1; /* local column index */
468 #endif
469         if (colb == -1) {
470           nrows = 0;
471         } else {
472           colb  = colb/bs;
473           row   = B_cj + B_ci[colb];
474           nrows = B_ci[colb+1] - B_ci[colb];
475         }
476         nrows_i += nrows;
477         /* loop over columns of B marking them in rowhit */
478         for (k=0; k<nrows; k++) {
479           /* set valaddrhit for part B */
480           spidx            = bs2*spidxB[B_ci[colb] + k];
481           valaddrhit[*row] = &B_val[spidx];
482           rowhit[*row++]   = colb + 1 + cend - cstart; /* local column index */
483         }
484       }
485     }
486     c->nrows[i] = nrows_i;
487 
488     for (j=0; j<m; j++) {
489       if (rowhit[j]) {
490         Jentry[nz].row     = j;              /* local row index */
491         Jentry[nz].col     = rowhit[j] - 1;  /* local column index */
492         Jentry[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
493         nz++;
494       }
495     }
496     ierr = PetscFree(cols);CHKERRQ(ierr);
497   }
498   ierr = ISColoringRestoreIS(iscoloring,&isa);CHKERRQ(ierr);
499 
500   ierr = PetscFree2(rowhit,valaddrhit);CHKERRQ(ierr);
501   if (isBAIJ) {
502     ierr = MatRestoreColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
503     ierr = MatRestoreColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
504     ierr = PetscMalloc(bs*mat->rmap->n*sizeof(PetscScalar),&c->dy);CHKERRQ(ierr);
505   } else {
506     ierr = MatRestoreColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
507     ierr = MatRestoreColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
508   }
509 
510   if (ctype == IS_COLORING_GHOSTED) {
511     ierr = ISLocalToGlobalMappingRestoreIndices(map,&ltog);CHKERRQ(ierr);
512   }
513   PetscFunctionReturn(0);
514 }
515