xref: /petsc/src/mat/impls/aij/mpi/fdmpiaij.c (revision 6ce644ab5e513ced4ec2e16b3fd30774bab3efb1)
1 
2 #include <../src/mat/impls/aij/mpi/mpiaij.h>
3 
4 extern PetscErrorCode MatCreateColmap_MPIAIJ_Private(Mat);
5 
6 #undef __FUNCT__
7 #define __FUNCT__ "MatFDColoringApply_MPIAIJ"
8 PetscErrorCode  MatFDColoringApply_MPIAIJ(Mat J,MatFDColoring coloring,Vec x1,MatStructure *flag,void *sctx)
9 {
10   PetscErrorCode (*f)(void*,Vec,Vec,void*) = (PetscErrorCode (*)(void*,Vec,Vec,void*))coloring->f;
11   PetscErrorCode ierr;
12   PetscInt       k,start,end,l,row,col,colb,**columnsforrow=coloring->columnsforrow,nz;
13   PetscScalar    dx,*y,*xx,*w3_array;
14   PetscScalar    *vscale_array;
15   PetscReal      epsilon = coloring->error_rel,umin = coloring->umin,unorm;
16   Vec            w1      = coloring->w1,w2=coloring->w2,w3;
17   void           *fctx   = coloring->fctx;
18   PetscBool      flg     = PETSC_FALSE;
19   PetscInt       ctype   = coloring->ctype,N,col_start=0,col_end=0;
20   Mat_MPIAIJ     *aij = (Mat_MPIAIJ*)J->data;
21   Mat            A=aij->A,B=aij->B;
22   Mat_SeqAIJ     *cspA=(Mat_SeqAIJ*)A->data, *cspB=(Mat_SeqAIJ*)B->data;
23   PetscMPIInt    rank;
24 
25   PetscFunctionBegin;
26   ierr = MPI_Comm_rank(PetscObjectComm((PetscObject)J), &rank);CHKERRQ(ierr);
27   ierr = MatSetUnfactored(J);CHKERRQ(ierr);
28   ierr = PetscOptionsGetBool(NULL,"-mat_fd_coloring_dont_rezero",&flg,NULL);CHKERRQ(ierr);
29   if (flg) {
30     ierr = PetscInfo(coloring,"Not calling MatZeroEntries()\n");CHKERRQ(ierr);
31   } else {
32     PetscBool assembled;
33     ierr = MatAssembled(J,&assembled);CHKERRQ(ierr);
34     if (assembled) {
35       ierr = MatZeroEntries(J);CHKERRQ(ierr);
36     }
37   }
38 
39   if (!coloring->vscale) { //ctype == IS_COLORING_GHOSTED ?
40     printf("[%d] create a non-ghostedcoloring->vscale\n",rank);
41     ierr = VecDuplicate(x1,&coloring->vscale);CHKERRQ(ierr);
42     //SETERRQ(PetscObjectComm((PetscObject)J),PETSC_ERR_ARG_NULL,"Null Object: coloring->scale");
43   }
44 
45   /* Set w1 = F(x1) */
46   if (!coloring->fset) {
47     ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
48     ierr = (*f)(sctx,x1,w1,fctx);CHKERRQ(ierr);
49     ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
50   } else {
51     coloring->fset = PETSC_FALSE;
52   }
53 
54   if (!coloring->w3) {
55     ierr = VecDuplicate(x1,&coloring->w3);CHKERRQ(ierr);
56     ierr = PetscLogObjectParent((PetscObject)coloring,(PetscObject)coloring->w3);CHKERRQ(ierr);
57   }
58   w3 = coloring->w3;
59 
60   /* Compute vscale = dx - the local scale factors, including ghost points */
61   ierr = VecGetOwnershipRange(w1,&start,&end);CHKERRQ(ierr); /* =J's row ownershipRange, used by ghosted x! */
62   //printf("[%d] start end %d %d\n",rank,start,end);
63 
64   if (coloring->htype[0] == 'w') {
65     /* tacky test; need to make systematic if we add other approaches to computing h*/
66     ierr = VecNorm(x1,NORM_2,&unorm);CHKERRQ(ierr);
67     dx = (1.0 + unorm)*epsilon;
68     ierr = VecSet(coloring->vscale,dx);CHKERRQ(ierr);
69   } else {
70     ierr = VecGetLocalSize(x1,&N);CHKERRQ(ierr);
71     ierr = VecGetArray(x1,&xx);CHKERRQ(ierr);
72     ierr = VecGetArray(coloring->vscale,&vscale_array);CHKERRQ(ierr);
73     if (ctype == IS_COLORING_GHOSTED) {
74       col_start = 0; col_end = N;
75     } else if (ctype == IS_COLORING_GLOBAL) {
76       xx           = xx - start;
77       vscale_array = vscale_array - start;
78       col_start    = start; col_end = N + start;
79     }
80     for (col=col_start; col<col_end; col++) {
81       /* Loop over each local column, vscale[col] = 1./(epsilon*dx[col]) */
82       dx = xx[col];
83       if (dx == (PetscScalar)0.0) dx = 1.0;
84       if (PetscAbsScalar(dx) < umin && PetscRealPart(dx) >= 0.0)     dx = umin;
85       else if (PetscRealPart(dx) < 0.0 && PetscAbsScalar(dx) < umin) dx = -umin;
86       dx               *= epsilon;
87       vscale_array[col] = (PetscScalar)dx;
88     }
89     if (ctype == IS_COLORING_GLOBAL) {
90       vscale_array = vscale_array + start;
91     }
92     ierr = VecRestoreArray(coloring->vscale,&vscale_array);CHKERRQ(ierr);
93   }
94 
95   if (ctype == IS_COLORING_GLOBAL) {
96     ierr = VecGhostUpdateBegin(coloring->vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
97     ierr = VecGhostUpdateEnd(coloring->vscale,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
98   }
99 
100   /* Loop over each color */
101   nz = 0;
102   ierr = VecGetArray(coloring->vscale,&vscale_array);CHKERRQ(ierr);
103   for (k=0; k<coloring->ncolors; k++) {
104     coloring->currentcolor = k;
105 
106     /*
107       Loop over each column associated with color
108       adding the perturbation to the vector w3 = x1 + dx.
109     */
110     ierr = VecCopy(x1,w3);CHKERRQ(ierr);
111     ierr = VecGetArray(w3,&w3_array);CHKERRQ(ierr);
112     if (ctype == IS_COLORING_GLOBAL) w3_array = w3_array - start;
113 
114     for (l=0; l<coloring->ncolumns[k]; l++) {
115       col = coloring->columns[k][l]; /* local column (in global index!) of the matrix we are probing for */
116       /* convert global col to local colb */
117       if (col >= start && col < end) {
118         colb = col-start;
119       } else {
120 #if defined(PETSC_USE_CTABLE)
121         ierr = PetscTableFind(aij->colmap,col+1,&colb);CHKERRQ(ierr);
122         colb--;
123 #else
124         colb = aij->colmap[col] - 1; /* local column index */
125 #endif
126         colb += end - start;
127       }
128       w3_array[col] += vscale_array[colb];
129     }
130     if (ctype == IS_COLORING_GLOBAL) w3_array = w3_array + start;
131     ierr = VecRestoreArray(w3,&w3_array);CHKERRQ(ierr);
132 
133     /*
134       Evaluate function at w3 = x1 + dx (here dx is a vector of perturbations)
135                            w2 = F(x1 + dx) - F(x1)
136     */
137     ierr = PetscLogEventBegin(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
138     ierr = (*f)(sctx,w3,w2,fctx);CHKERRQ(ierr);
139     ierr = PetscLogEventEnd(MAT_FDColoringFunction,0,0,0,0);CHKERRQ(ierr);
140     ierr = VecAXPY(w2,-1.0,w1);CHKERRQ(ierr);
141 
142     /* Loop over rows of vector, putting results into Jacobian matrix */
143     ierr = VecGetArray(w2,&y);CHKERRQ(ierr);
144     for (l=0; l<coloring->nrows[k]; l++) {
145       row                        = coloring->rows[k][l];   /* local row index */
146       *(coloring->valaddr[nz++]) = y[row]/vscale_array[columnsforrow[k][l]];
147     }
148     ierr = VecRestoreArray(w2,&y);CHKERRQ(ierr);
149   } /* endof for each color */
150   if (nz != cspA->nz + cspB->nz) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"nz %d != mat->nz %d\n",nz,cspA->nz+cspB->nz);
151   ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
152   ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
153 
154   if (ctype == IS_COLORING_GLOBAL) xx = xx + start;
155   ierr = VecRestoreArray(coloring->vscale,&vscale_array);CHKERRQ(ierr);
156   ierr = VecRestoreArray(x1,&xx);CHKERRQ(ierr);
157 
158   coloring->currentcolor = -1;
159   PetscFunctionReturn(0);
160 }
161 
162 #undef __FUNCT__
163 #define __FUNCT__ "MatFDColoringCreate_MPIAIJ_new"
164 PetscErrorCode MatFDColoringCreate_MPIAIJ_new(Mat mat,ISColoring iscoloring,MatFDColoring c)
165 {
166   Mat_MPIAIJ             *aij = (Mat_MPIAIJ*)mat->data;
167   PetscErrorCode         ierr;
168   PetscMPIInt            size,*ncolsonproc,*disp,nn;
169   PetscInt               i,n,nrows,j,k,m,ncols,col;
170   const PetscInt         *is,*A_ci,*A_cj,*B_ci,*B_cj,*row = NULL,*ltog=NULL;
171   PetscInt               nis = iscoloring->n,nctot,*cols;
172   PetscInt               *rowhit,cstart,cend,colb;
173   IS                     *isa;
174   ISLocalToGlobalMapping map = mat->cmap->mapping;
175   PetscInt               ctype=c->ctype;
176   Mat                    A=aij->A,B=aij->B;
177   Mat_SeqAIJ             *cspA=(Mat_SeqAIJ*)A->data, *cspB=(Mat_SeqAIJ*)B->data;
178   PetscScalar            *A_val=cspA->a,*B_val=cspB->a;
179   PetscInt               spidx;
180   PetscInt               *spidxA,*spidxB,nz;
181   PetscScalar            **valaddrhit;
182 
183   PetscFunctionBegin;
184   if (ctype == IS_COLORING_GHOSTED) {
185     if(!map) SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_INCOMP,"When using ghosted differencing matrix must have local to global mapping provided with MatSetLocalToGlobalMapping");
186     ierr = ISLocalToGlobalMappingGetIndices(map,&ltog);CHKERRQ(ierr);
187   }
188 
189   m         = mat->rmap->n;
190   cstart    = mat->cmap->rstart;
191   cend      = mat->cmap->rend;
192   c->M      = mat->rmap->N;         /* set the global rows and columns and local rows */
193   c->N      = mat->cmap->N;
194   c->m      = m;
195   c->rstart = mat->rmap->rstart;
196 
197   c->ncolors = nis;
198   ierr       = PetscMalloc(nis*sizeof(PetscInt),&c->ncolumns);CHKERRQ(ierr);
199   ierr       = PetscMalloc(nis*sizeof(PetscInt*),&c->columns);CHKERRQ(ierr);
200   ierr       = PetscMalloc(nis*sizeof(PetscInt),&c->nrows);CHKERRQ(ierr);
201   ierr       = PetscMalloc(nis*sizeof(PetscInt*),&c->rows);CHKERRQ(ierr);
202   ierr       = PetscMalloc(nis*sizeof(PetscInt*),&c->columnsforrow);CHKERRQ(ierr);
203   ierr       = PetscLogObjectMemory((PetscObject)c,5*nis*sizeof(PetscInt));CHKERRQ(ierr);
204 
205   nz         = cspA->nz + cspB->nz; /* total nonzero entries of mat */
206   ierr       = PetscMalloc(nz*sizeof(PetscScalar*),&c->valaddr);CHKERRQ(ierr);
207   ierr       = PetscLogObjectMemory((PetscObject)c,nz*sizeof(PetscScalar*));CHKERRQ(ierr);
208 
209   /* Allow access to data structures of local part of matrix
210    - creates aij->colmap which maps global column number to local number in part B */
211   if (!aij->colmap) {
212     ierr = MatCreateColmap_MPIAIJ_Private(mat);CHKERRQ(ierr);
213   }
214   ierr = MatGetColumnIJ_SeqAIJ_Color(aij->A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
215   ierr = MatGetColumnIJ_SeqAIJ_Color(aij->B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
216 
217   ierr = PetscMalloc2(m+1,PetscInt,&rowhit,m+1,PetscScalar*,&valaddrhit);CHKERRQ(ierr);
218   nz = 0;
219   ierr = ISColoringGetIS(iscoloring,PETSC_IGNORE,&isa);CHKERRQ(ierr);
220   for (i=0; i<nis; i++) { /* for each local color */
221     ierr = ISGetLocalSize(isa[i],&n);CHKERRQ(ierr);
222     ierr = ISGetIndices(isa[i],&is);CHKERRQ(ierr);
223 
224     c->ncolumns[i] = n; /* local number of columns of this color on this process */
225     if (n) {
226       ierr = PetscMalloc(n*sizeof(PetscInt),&c->columns[i]);CHKERRQ(ierr);
227       ierr = PetscLogObjectMemory((PetscObject)c,n*sizeof(PetscInt));CHKERRQ(ierr);
228       ierr = PetscMemcpy(c->columns[i],is,n*sizeof(PetscInt));CHKERRQ(ierr);
229     } else {
230       c->columns[i] = 0;
231     }
232 
233     if (ctype == IS_COLORING_GLOBAL) {
234       /* Determine nctot, the total (parallel) number of columns of this color */
235       ierr = MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size);CHKERRQ(ierr);
236       ierr = PetscMalloc2(size,PetscMPIInt,&ncolsonproc,size,PetscMPIInt,&disp);CHKERRQ(ierr);
237 
238       /* ncolsonproc[j]: local ncolumns on proc[j] of this color */
239       ierr  = PetscMPIIntCast(n,&nn);CHKERRQ(ierr);
240       ierr  = MPI_Allgather(&nn,1,MPI_INT,ncolsonproc,1,MPI_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr);
241       nctot = 0; for (j=0; j<size; j++) nctot += ncolsonproc[j];
242       if (!nctot) {
243         ierr = PetscInfo(mat,"Coloring of matrix has some unneeded colors with no corresponding rows\n");CHKERRQ(ierr);
244       }
245 
246       disp[0] = 0;
247       for (j=1; j<size; j++) {
248         disp[j] = disp[j-1] + ncolsonproc[j-1];
249       }
250 
251       /* Get cols, the complete list of columns for this color on each process */
252       ierr = PetscMalloc((nctot+1)*sizeof(PetscInt),&cols);CHKERRQ(ierr);
253       ierr = MPI_Allgatherv((void*)is,n,MPIU_INT,cols,ncolsonproc,disp,MPIU_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr);
254       ierr = PetscFree2(ncolsonproc,disp);CHKERRQ(ierr);
255     } else if (ctype == IS_COLORING_GHOSTED) {
256       /* Determine local number of columns of this color on this process, including ghost points */
257       nctot = n;
258       ierr  = PetscMalloc((nctot+1)*sizeof(PetscInt),&cols);CHKERRQ(ierr);
259       ierr  = PetscMemcpy(cols,is,n*sizeof(PetscInt));CHKERRQ(ierr);
260     } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not provided for this MatFDColoring type");
261 
262     /* Mark all rows affect by these columns */
263     ierr = PetscMemzero(rowhit,m*sizeof(PetscInt));CHKERRQ(ierr);
264     for (j=0; j<nctot; j++) { /* loop over columns*/
265       if (ctype == IS_COLORING_GHOSTED) {
266         col = ltog[cols[j]];
267       } else {
268         col = cols[j];
269       }
270       if (col >= cstart && col < cend) { /* column is in diagonal block of matrix A */
271         row = A_cj + A_ci[col-cstart];
272         nrows = A_ci[col-cstart+1] - A_ci[col-cstart];
273         /* loop over columns of A marking them in rowhit */
274         for (k=0; k<nrows; k++) {
275           /* set valaddrhit for part A */
276           spidx            = spidxA[A_ci[col-cstart] + k];
277           valaddrhit[*row] = &A_val[spidx];
278           rowhit[*row++]   = col - cstart + 1; /* local column index */
279         }
280       } else { /* column is in off-diagonal block of matrix B */
281 #if defined(PETSC_USE_CTABLE)
282         ierr = PetscTableFind(aij->colmap,col+1,&colb);CHKERRQ(ierr);
283         colb--;
284 #else
285         colb = aij->colmap[col] - 1; /* local column index */
286 #endif
287         if (colb == -1) {
288           nrows = 0;
289         } else {
290           row   = B_cj + B_ci[colb];
291           nrows = B_ci[colb+1] - B_ci[colb];
292         }
293         /* loop over columns of B marking them in rowhit */
294         for (k=0; k<nrows; k++) {
295           /* set valaddrhit for part B */
296           spidx            = spidxB[B_ci[colb] + k];
297           valaddrhit[*row] = &B_val[spidx];
298           rowhit[*row++]   = colb + 1 + cend - cstart; /* local column index */
299         }
300       }
301     }
302 
303     /* count the number of hits */
304     nrows = 0;
305     for (j=0; j<m; j++) {
306       if (rowhit[j]) nrows++;
307     }
308     c->nrows[i] = nrows;
309     ierr        = PetscMalloc((nrows+1)*sizeof(PetscInt),&c->rows[i]);CHKERRQ(ierr);
310     ierr        = PetscMalloc((nrows+1)*sizeof(PetscInt),&c->columnsforrow[i]);CHKERRQ(ierr);
311     ierr        = PetscLogObjectMemory((PetscObject)c,2*(nrows+1)*sizeof(PetscInt));CHKERRQ(ierr);
312     nrows       = 0;
313     for (j=0; j<m; j++) {
314       if (rowhit[j]) {
315         c->rows[i][nrows]          = j;              /* local row index */
316         c->columnsforrow[i][nrows] = rowhit[j] - 1;  /* local column index */
317         c->valaddr[nz++]           = valaddrhit[j];  /* address of mat value for this entry */
318         nrows++;
319       }
320     }
321     ierr = PetscFree(cols);CHKERRQ(ierr);
322   }
323   if (nz != cspA->nz + cspB->nz) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"nz %d != mat->nz %d\n",nz,cspA->nz+cspB->nz);
324 
325   /* create vscale for storing dx */
326   if (ctype == IS_COLORING_GLOBAL) {
327     ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),m,PETSC_DETERMINE,B->cmap->n,aij->garray,&c->vscale);CHKERRQ(ierr);
328   }
329 
330   ierr = ISColoringRestoreIS(iscoloring,&isa);CHKERRQ(ierr);
331 
332   ierr = PetscFree2(rowhit,valaddrhit);CHKERRQ(ierr);
333   ierr = MatRestoreColumnIJ_SeqAIJ_Color(aij->A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
334   ierr = MatRestoreColumnIJ_SeqAIJ_Color(aij->B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
335   if (ctype == IS_COLORING_GHOSTED) {
336     ierr = ISLocalToGlobalMappingRestoreIndices(map,&ltog);CHKERRQ(ierr);
337   }
338 
339   mat->ops->fdcoloringapply = MatFDColoringApply_MPIAIJ;
340   PetscFunctionReturn(0);
341 }
342 
343 /*------------------------------------------------------*/
344 #undef __FUNCT__
345 #define __FUNCT__ "MatFDColoringCreate_MPIAIJ"
346 PetscErrorCode MatFDColoringCreate_MPIAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
347 {
348   Mat_MPIAIJ             *aij = (Mat_MPIAIJ*)mat->data;
349   PetscErrorCode         ierr;
350   PetscMPIInt            size,*ncolsonproc,*disp,nn;
351   PetscInt               i,n,nrows,j,k,m,ncols,col;
352   const PetscInt         *is,*A_ci,*A_cj,*B_ci,*B_cj,*rows = 0,*ltog;
353   PetscInt               nis = iscoloring->n,nctot,*cols;
354   PetscInt               *rowhit,M,cstart,cend,colb;
355   PetscInt               *columnsforrow,l;
356   IS                     *isa;
357   PetscBool              done,flg;
358   ISLocalToGlobalMapping map   = mat->cmap->mapping;
359   PetscInt               ctype=c->ctype;
360   PetscBool              new_impl=PETSC_FALSE;
361 
362   PetscFunctionBegin;
363   ierr = PetscOptionsName("-new","using new impls","",&new_impl);CHKERRQ(ierr);
364   if (new_impl){
365     ierr =  MatFDColoringCreate_MPIAIJ_new(mat,iscoloring,c);CHKERRQ(ierr);
366     PetscFunctionReturn(0);
367   }
368   if (ctype == IS_COLORING_GHOSTED && !map) SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_INCOMP,"When using ghosted differencing matrix must have local to global mapping provided with MatSetLocalToGlobalMapping");
369 
370   if (map) {ierr = ISLocalToGlobalMappingGetIndices(map,&ltog);CHKERRQ(ierr);}
371   else     ltog = NULL;
372   ierr = ISColoringGetIS(iscoloring,PETSC_IGNORE,&isa);CHKERRQ(ierr);
373 
374   M         = mat->rmap->n;
375   cstart    = mat->cmap->rstart;
376   cend      = mat->cmap->rend;
377   c->M      = mat->rmap->N;         /* set the global rows and columns and local rows */
378   c->N      = mat->cmap->N;
379   c->m      = mat->rmap->n;
380   c->rstart = mat->rmap->rstart;
381 
382   c->ncolors = nis;
383   ierr       = PetscMalloc(nis*sizeof(PetscInt),&c->ncolumns);CHKERRQ(ierr);
384   ierr       = PetscMalloc(nis*sizeof(PetscInt*),&c->columns);CHKERRQ(ierr);
385   ierr       = PetscMalloc(nis*sizeof(PetscInt),&c->nrows);CHKERRQ(ierr);
386   ierr       = PetscMalloc(nis*sizeof(PetscInt*),&c->rows);CHKERRQ(ierr);
387   ierr       = PetscMalloc(nis*sizeof(PetscInt*),&c->columnsforrow);CHKERRQ(ierr);
388   ierr       = PetscLogObjectMemory((PetscObject)c,5*nis*sizeof(PetscInt));CHKERRQ(ierr);
389 
390   /* Allow access to data structures of local part of matrix */
391   if (!aij->colmap) {
392     ierr = MatCreateColmap_MPIAIJ_Private(mat);CHKERRQ(ierr);
393   }
394   ierr = MatGetColumnIJ(aij->A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&done);CHKERRQ(ierr);
395   ierr = MatGetColumnIJ(aij->B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&done);CHKERRQ(ierr);
396 
397   ierr = PetscMalloc((M+1)*sizeof(PetscInt),&rowhit);CHKERRQ(ierr);
398   ierr = PetscMalloc((M+1)*sizeof(PetscInt),&columnsforrow);CHKERRQ(ierr);
399 
400   for (i=0; i<nis; i++) {
401     ierr = ISGetLocalSize(isa[i],&n);CHKERRQ(ierr);
402     ierr = ISGetIndices(isa[i],&is);CHKERRQ(ierr);
403 
404     c->ncolumns[i] = n; /* local number of columns of this color on this process */
405     if (n) {
406       ierr = PetscMalloc(n*sizeof(PetscInt),&c->columns[i]);CHKERRQ(ierr);
407       ierr = PetscLogObjectMemory((PetscObject)c,n*sizeof(PetscInt));CHKERRQ(ierr);
408       ierr = PetscMemcpy(c->columns[i],is,n*sizeof(PetscInt));CHKERRQ(ierr);
409     } else {
410       c->columns[i] = 0;
411     }
412 
413     if (ctype == IS_COLORING_GLOBAL) {
414       /* Determine the total (parallel) number of columns of this color */
415       ierr = MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size);CHKERRQ(ierr);
416       ierr = PetscMalloc2(size,PetscMPIInt,&ncolsonproc,size,PetscMPIInt,&disp);CHKERRQ(ierr);
417 
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 complete list of columns for color on each processor */
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     /*
442        Mark all rows affect by these columns
443     */
444     /* Temporary option to allow for debugging/testing */
445     flg  = PETSC_FALSE;
446     ierr = PetscOptionsGetBool(NULL,"-matfdcoloring_slow",&flg,NULL);CHKERRQ(ierr);
447     if (!flg) { /*-----------------------------------------------------------------------------*/
448       /* crude, fast version */
449       ierr = PetscMemzero(rowhit,M*sizeof(PetscInt));CHKERRQ(ierr);
450       /* loop over columns*/
451       for (j=0; j<nctot; j++) {
452         if (ctype == IS_COLORING_GHOSTED) {
453           col = ltog[cols[j]];
454         } else {
455           col = cols[j];
456         }
457         if (col >= cstart && col < cend) {
458           /* column is in diagonal block of matrix */
459           rows = A_cj + A_ci[col-cstart];
460           m    = A_ci[col-cstart+1] - A_ci[col-cstart];
461         } else {
462 #if defined(PETSC_USE_CTABLE)
463           ierr = PetscTableFind(aij->colmap,col+1,&colb);CHKERRQ(ierr);
464           colb--;
465 #else
466           colb = aij->colmap[col] - 1;
467 #endif
468           if (colb == -1) {
469             m = 0;
470           } else {
471             rows = B_cj + B_ci[colb];
472             m    = B_ci[colb+1] - B_ci[colb];
473           }
474         }
475         /* loop over columns marking them in rowhit */
476         for (k=0; k<m; k++) {
477           rowhit[*rows++] = col + 1;
478         }
479       }
480 
481       /* count the number of hits */
482       nrows = 0;
483       for (j=0; j<M; j++) {
484         if (rowhit[j]) nrows++;
485       }
486       c->nrows[i] = nrows;
487       ierr        = PetscMalloc((nrows+1)*sizeof(PetscInt),&c->rows[i]);CHKERRQ(ierr);
488       ierr        = PetscMalloc((nrows+1)*sizeof(PetscInt),&c->columnsforrow[i]);CHKERRQ(ierr);
489       ierr        = PetscLogObjectMemory((PetscObject)c,2*(nrows+1)*sizeof(PetscInt));CHKERRQ(ierr);
490       nrows       = 0;
491       for (j=0; j<M; j++) {
492         if (rowhit[j]) {
493           c->rows[i][nrows]          = j;              /* local row index */
494           c->columnsforrow[i][nrows] = rowhit[j] - 1;  /* global column index */
495           nrows++;
496         }
497       }
498     } else { /*-------------------------------------------------------------------------------*/
499       /* slow version, using rowhit as a linked list */
500       PetscInt currentcol,fm,mfm;
501       rowhit[M] = M;
502       nrows     = 0;
503       /* loop over columns*/
504       for (j=0; j<nctot; j++) {
505         if (ctype == IS_COLORING_GHOSTED) {
506           col = ltog[cols[j]];
507         } else {
508           col = cols[j];
509         }
510         if (col >= cstart && col < cend) {
511           /* column is in diagonal block of matrix */
512           rows = A_cj + A_ci[col-cstart];
513           m    = A_ci[col-cstart+1] - A_ci[col-cstart];
514         } else {
515 #if defined(PETSC_USE_CTABLE)
516           ierr = PetscTableFind(aij->colmap,col+1,&colb);CHKERRQ(ierr);
517           colb--;
518 #else
519           colb = aij->colmap[col] - 1;
520 #endif
521           if (colb == -1) {
522             m = 0;
523           } else {
524             rows = B_cj + B_ci[colb];
525             m    = B_ci[colb+1] - B_ci[colb];
526           }
527         }
528 
529         /* loop over columns marking them in rowhit */
530         fm = M;    /* fm points to first entry in linked list */
531         for (k=0; k<m; k++) {
532           currentcol = *rows++;
533           /* is it already in the list? */
534           do {
535             mfm = fm;
536             fm  = rowhit[fm];
537           } while (fm < currentcol);
538           /* not in list so add it */
539           if (fm != currentcol) {
540             nrows++;
541             columnsforrow[currentcol] = col;
542             /* next three lines insert new entry into linked list */
543             rowhit[mfm]        = currentcol;
544             rowhit[currentcol] = fm;
545             fm                 = currentcol;
546             /* fm points to present position in list since we know the columns are sorted */
547           } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Invalid coloring of matrix detected");
548         }
549       }
550       c->nrows[i] = nrows;
551 
552       ierr = PetscMalloc((nrows+1)*sizeof(PetscInt),&c->rows[i]);CHKERRQ(ierr);
553       ierr = PetscMalloc((nrows+1)*sizeof(PetscInt),&c->columnsforrow[i]);CHKERRQ(ierr);
554       ierr = PetscLogObjectMemory((PetscObject)c,(nrows+1)*sizeof(PetscInt));CHKERRQ(ierr);
555       /* now store the linked list of rows into c->rows[i] */
556       nrows = 0;
557       fm    = rowhit[M];
558       do {
559         c->rows[i][nrows]            = fm;
560         c->columnsforrow[i][nrows++] = columnsforrow[fm];
561         fm                           = rowhit[fm];
562       } while (fm < M);
563     } /* ---------------------------------------------------------------------------------------*/
564     ierr = PetscFree(cols);CHKERRQ(ierr);
565   }
566 
567   /* Optimize by adding the vscale, and scaleforrow[][] fields */
568   /*
569        vscale will contain the "diagonal" on processor scalings followed by the off processor
570   */
571   if (ctype == IS_COLORING_GLOBAL) {
572     ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),aij->A->rmap->n,PETSC_DETERMINE,aij->B->cmap->n,aij->garray,&c->vscale);CHKERRQ(ierr);
573     ierr = PetscMalloc(c->ncolors*sizeof(PetscInt*),&c->vscaleforrow);CHKERRQ(ierr);
574     for (k=0; k<c->ncolors; k++) {
575       ierr = PetscMalloc((c->nrows[k]+1)*sizeof(PetscInt),&c->vscaleforrow[k]);CHKERRQ(ierr);
576       for (l=0; l<c->nrows[k]; l++) {
577         col = c->columnsforrow[k][l];
578         if (col >= cstart && col < cend) {
579           /* column is in diagonal block of matrix */
580           colb = col - cstart;
581         } else {
582           /* column  is in "off-processor" part */
583 #if defined(PETSC_USE_CTABLE)
584           ierr = PetscTableFind(aij->colmap,col+1,&colb);CHKERRQ(ierr);
585           colb--;
586 #else
587           colb = aij->colmap[col] - 1;
588 #endif
589           colb += cend - cstart;
590         }
591         c->vscaleforrow[k][l] = colb;
592       }
593     }
594   } else if (ctype == IS_COLORING_GHOSTED) {
595     /* Get gtol mapping */
596     PetscInt N = mat->cmap->N,nlocal,*gtol;
597     ierr = PetscMalloc((N+1)*sizeof(PetscInt),&gtol);CHKERRQ(ierr);
598     for (i=0; i<N; i++) gtol[i] = -1;
599     ierr = ISLocalToGlobalMappingGetSize(map,&nlocal);CHKERRQ(ierr);
600     for (i=0; i<nlocal; i++) gtol[ltog[i]] = i;
601 
602     c->vscale = 0; /* will be created in MatFDColoringApply() */
603     ierr      = PetscMalloc(c->ncolors*sizeof(PetscInt*),&c->vscaleforrow);CHKERRQ(ierr);
604     for (k=0; k<c->ncolors; k++) {
605       ierr = PetscMalloc((c->nrows[k]+1)*sizeof(PetscInt),&c->vscaleforrow[k]);CHKERRQ(ierr);
606       for (l=0; l<c->nrows[k]; l++) {
607         col = c->columnsforrow[k][l];      /* global column index */
608         c->vscaleforrow[k][l] = gtol[col]; /* local column index */
609       }
610     }
611     ierr = PetscFree(gtol);CHKERRQ(ierr);
612   }
613   ierr = ISColoringRestoreIS(iscoloring,&isa);CHKERRQ(ierr);
614 
615   ierr = PetscFree(rowhit);CHKERRQ(ierr);
616   ierr = PetscFree(columnsforrow);CHKERRQ(ierr);
617   ierr = MatRestoreColumnIJ(aij->A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&done);CHKERRQ(ierr);
618   ierr = MatRestoreColumnIJ(aij->B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&done);CHKERRQ(ierr);
619   if (map) {ierr = ISLocalToGlobalMappingRestoreIndices(map,&ltog);CHKERRQ(ierr);}
620   PetscFunctionReturn(0);
621 }
622 
623 
624 
625 
626 
627 
628