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