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