xref: /petsc/src/mat/impls/aij/mpi/fdmpiaij.c (revision fc8a9adeb7fcdc98711d755fa2dc544ddccf0f3e) ! 
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;
372   IS                     *isa;
373   ISLocalToGlobalMapping map=mat->cmap->mapping;
374   PetscInt               ctype=c->ctype,*spidxA,*spidxB,nz,bs,bs2,spidx;
375   Mat                    A,B;
376   PetscScalar            *A_val,*B_val,**valaddrhit;
377   MatEntry               *Jentry;
378   MatEntry2              *Jentry2;
379   PetscBool              isBAIJ,isSELL;
380   PetscInt               bcols=c->bcols;
381 #if defined(PETSC_USE_CTABLE)
382   PetscTable             colmap=NULL;
383 #else
384   PetscInt               *colmap=NULL;     /* local col number of off-diag col */
385 #endif
386 
387   PetscFunctionBegin;
388   if (ctype == IS_COLORING_LOCAL) {
389     if (!map) SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_INCOMP,"When using ghosted differencing matrix must have local to global mapping provided with MatSetLocalToGlobalMapping");
390     ierr = ISLocalToGlobalMappingGetIndices(map,&ltog);CHKERRQ(ierr);
391   }
392 
393   ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
394   ierr = PetscObjectTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);CHKERRQ(ierr);
395   ierr = PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL);CHKERRQ(ierr);
396   if (isBAIJ) {
397     Mat_MPIBAIJ *baij=(Mat_MPIBAIJ*)mat->data;
398     Mat_SeqBAIJ *spA,*spB;
399     A = baij->A;  spA = (Mat_SeqBAIJ*)A->data; A_val = spA->a;
400     B = baij->B;  spB = (Mat_SeqBAIJ*)B->data; B_val = spB->a;
401     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
402     if (!baij->colmap) {
403       ierr = MatCreateColmap_MPIBAIJ_Private(mat);CHKERRQ(ierr);
404     }
405     colmap = baij->colmap;
406     ierr = MatGetColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
407     ierr = MatGetColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
408 
409     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') {  /* create vscale for storing dx */
410       PetscInt    *garray;
411       ierr = PetscMalloc1(B->cmap->n,&garray);CHKERRQ(ierr);
412       for (i=0; i<baij->B->cmap->n/bs; i++) {
413         for (j=0; j<bs; j++) {
414           garray[i*bs+j] = bs*baij->garray[i]+j;
415         }
416       }
417       ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,garray,&c->vscale);CHKERRQ(ierr);
418       ierr = VecPinToCPU(c->vscale,PETSC_TRUE);CHKERRQ(ierr);
419       ierr = PetscFree(garray);CHKERRQ(ierr);
420     }
421   } else if (isSELL) {
422     Mat_MPISELL *sell=(Mat_MPISELL*)mat->data;
423     Mat_SeqSELL *spA,*spB;
424     A = sell->A;  spA = (Mat_SeqSELL*)A->data; A_val = spA->val;
425     B = sell->B;  spB = (Mat_SeqSELL*)B->data; B_val = spB->val;
426     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
427     if (!sell->colmap) {
428       /* Allow access to data structures of local part of matrix
429        - creates aij->colmap which maps global column number to local number in part B */
430       ierr = MatCreateColmap_MPISELL_Private(mat);CHKERRQ(ierr);
431     }
432     colmap = sell->colmap;
433     ierr = MatGetColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
434     ierr = MatGetColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
435 
436     bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */
437 
438     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
439       ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,sell->garray,&c->vscale);CHKERRQ(ierr);
440       ierr = VecPinToCPU(c->vscale,PETSC_TRUE);CHKERRQ(ierr);
441     }
442   } else {
443     Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
444     Mat_SeqAIJ *spA,*spB;
445     A = aij->A;  spA = (Mat_SeqAIJ*)A->data; A_val = spA->a;
446     B = aij->B;  spB = (Mat_SeqAIJ*)B->data; B_val = spB->a;
447     nz = spA->nz + spB->nz; /* total nonzero entries of mat */
448     if (!aij->colmap) {
449       /* Allow access to data structures of local part of matrix
450        - creates aij->colmap which maps global column number to local number in part B */
451       ierr = MatCreateColmap_MPIAIJ_Private(mat);CHKERRQ(ierr);
452     }
453     colmap = aij->colmap;
454     ierr = MatGetColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
455     ierr = MatGetColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
456 
457     bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */
458 
459     if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
460       ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,aij->garray,&c->vscale);CHKERRQ(ierr);
461       ierr = VecPinToCPU(c->vscale,PETSC_TRUE);CHKERRQ(ierr);
462     }
463   }
464 
465   m         = mat->rmap->n/bs;
466   cstart    = mat->cmap->rstart/bs;
467   cend      = mat->cmap->rend/bs;
468 
469   ierr       = PetscMalloc1(nis,&c->ncolumns);CHKERRQ(ierr);
470   ierr       = PetscMalloc1(nis,&c->columns);CHKERRQ(ierr);
471   ierr       = PetscCalloc1(nis,&c->nrows);CHKERRQ(ierr);
472   ierr       = PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));CHKERRQ(ierr);
473 
474   if (c->htype[0] == 'd') {
475     ierr       = PetscMalloc1(nz,&Jentry);CHKERRQ(ierr);
476     ierr       = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));CHKERRQ(ierr);
477     c->matentry = Jentry;
478   } else if (c->htype[0] == 'w') {
479     ierr       = PetscMalloc1(nz,&Jentry2);CHKERRQ(ierr);
480     ierr       = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry2));CHKERRQ(ierr);
481     c->matentry2 = Jentry2;
482   } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"htype is not supported");
483 
484   ierr = PetscMalloc2(m+1,&rowhit,m+1,&valaddrhit);CHKERRQ(ierr);
485   nz = 0;
486   ierr = ISColoringGetIS(iscoloring,PETSC_IGNORE,&isa);CHKERRQ(ierr);
487   for (i=0; i<nis; i++) { /* for each local color */
488     ierr = ISGetLocalSize(isa[i],&n);CHKERRQ(ierr);
489     ierr = ISGetIndices(isa[i],&is);CHKERRQ(ierr);
490 
491     c->ncolumns[i] = n; /* local number of columns of this color on this process */
492     if (n) {
493       ierr = PetscMalloc1(n,&c->columns[i]);CHKERRQ(ierr);
494       ierr = PetscLogObjectMemory((PetscObject)c,n*sizeof(PetscInt));CHKERRQ(ierr);
495       ierr = PetscMemcpy(c->columns[i],is,n*sizeof(PetscInt));CHKERRQ(ierr);
496     } else {
497       c->columns[i] = 0;
498     }
499 
500     if (ctype == IS_COLORING_GLOBAL) {
501       /* Determine nctot, the total (parallel) number of columns of this color */
502       ierr = MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size);CHKERRQ(ierr);
503       ierr = PetscMalloc2(size,&ncolsonproc,size,&disp);CHKERRQ(ierr);
504 
505       /* ncolsonproc[j]: local ncolumns on proc[j] of this color */
506       ierr  = PetscMPIIntCast(n,&nn);CHKERRQ(ierr);
507       ierr  = MPI_Allgather(&nn,1,MPI_INT,ncolsonproc,1,MPI_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr);
508       nctot = 0; for (j=0; j<size; j++) nctot += ncolsonproc[j];
509       if (!nctot) {
510         ierr = PetscInfo(mat,"Coloring of matrix has some unneeded colors with no corresponding rows\n");CHKERRQ(ierr);
511       }
512 
513       disp[0] = 0;
514       for (j=1; j<size; j++) {
515         disp[j] = disp[j-1] + ncolsonproc[j-1];
516       }
517 
518       /* Get cols, the complete list of columns for this color on each process */
519       ierr = PetscMalloc1(nctot+1,&cols);CHKERRQ(ierr);
520       ierr = MPI_Allgatherv((void*)is,n,MPIU_INT,cols,ncolsonproc,disp,MPIU_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr);
521       ierr = PetscFree2(ncolsonproc,disp);CHKERRQ(ierr);
522     } else if (ctype == IS_COLORING_LOCAL) {
523       /* Determine local number of columns of this color on this process, including ghost points */
524       nctot = n;
525       ierr  = PetscMalloc1(nctot+1,&cols);CHKERRQ(ierr);
526       ierr  = PetscMemcpy(cols,is,n*sizeof(PetscInt));CHKERRQ(ierr);
527     } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not provided for this MatFDColoring type");
528 
529     /* Mark all rows affect by these columns */
530     ierr    = PetscMemzero(rowhit,m*sizeof(PetscInt));CHKERRQ(ierr);
531     bs2     = bs*bs;
532     nrows_i = 0;
533     for (j=0; j<nctot; j++) { /* loop over columns*/
534       if (ctype == IS_COLORING_LOCAL) {
535         col = ltog[cols[j]];
536       } else {
537         col = cols[j];
538       }
539       if (col >= cstart && col < cend) { /* column is in A, diagonal block of mat */
540         row      = A_cj + A_ci[col-cstart];
541         nrows    = A_ci[col-cstart+1] - A_ci[col-cstart];
542         nrows_i += nrows;
543         /* loop over columns of A marking them in rowhit */
544         for (k=0; k<nrows; k++) {
545           /* set valaddrhit for part A */
546           spidx            = bs2*spidxA[A_ci[col-cstart] + k];
547           valaddrhit[*row] = &A_val[spidx];
548           rowhit[*row++]   = col - cstart + 1; /* local column index */
549         }
550       } else { /* column is in B, off-diagonal block of mat */
551 #if defined(PETSC_USE_CTABLE)
552         ierr = PetscTableFind(colmap,col+1,&colb);CHKERRQ(ierr);
553         colb--;
554 #else
555         colb = colmap[col] - 1; /* local column index */
556 #endif
557         if (colb == -1) {
558           nrows = 0;
559         } else {
560           colb  = colb/bs;
561           row   = B_cj + B_ci[colb];
562           nrows = B_ci[colb+1] - B_ci[colb];
563         }
564         nrows_i += nrows;
565         /* loop over columns of B marking them in rowhit */
566         for (k=0; k<nrows; k++) {
567           /* set valaddrhit for part B */
568           spidx            = bs2*spidxB[B_ci[colb] + k];
569           valaddrhit[*row] = &B_val[spidx];
570           rowhit[*row++]   = colb + 1 + cend - cstart; /* local column index */
571         }
572       }
573     }
574     c->nrows[i] = nrows_i;
575 
576     if (c->htype[0] == 'd') {
577       for (j=0; j<m; j++) {
578         if (rowhit[j]) {
579           Jentry[nz].row     = j;              /* local row index */
580           Jentry[nz].col     = rowhit[j] - 1;  /* local column index */
581           Jentry[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
582           nz++;
583         }
584       }
585     } else { /* c->htype == 'wp' */
586       for (j=0; j<m; j++) {
587         if (rowhit[j]) {
588           Jentry2[nz].row     = j;              /* local row index */
589           Jentry2[nz].valaddr = valaddrhit[j];  /* address of mat value for this entry */
590           nz++;
591         }
592       }
593     }
594     ierr = PetscFree(cols);CHKERRQ(ierr);
595   }
596 
597   if (bcols > 1) { /* reorder Jentry for faster MatFDColoringApply() */
598     ierr = MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz);CHKERRQ(ierr);
599   }
600 
601   if (isBAIJ) {
602     ierr = MatRestoreColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
603     ierr = MatRestoreColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
604     ierr = PetscMalloc1(bs*mat->rmap->n,&c->dy);CHKERRQ(ierr);
605   } else if (isSELL) {
606     ierr = MatRestoreColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
607     ierr = MatRestoreColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
608   }else {
609     ierr = MatRestoreColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
610     ierr = MatRestoreColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
611   }
612 
613   ierr = ISColoringRestoreIS(iscoloring,&isa);CHKERRQ(ierr);
614   ierr = PetscFree2(rowhit,valaddrhit);CHKERRQ(ierr);
615 
616   if (ctype == IS_COLORING_LOCAL) {
617     ierr = ISLocalToGlobalMappingRestoreIndices(map,&ltog);CHKERRQ(ierr);
618   }
619   ierr = PetscInfo3(c,"ncolors %D, brows %D and bcols %D are used.\n",c->ncolors,c->brows,c->bcols);CHKERRQ(ierr);
620   PetscFunctionReturn(0);
621 }
622 
623 PetscErrorCode MatFDColoringCreate_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
624 {
625   PetscErrorCode ierr;
626   PetscInt       bs,nis=iscoloring->n,m=mat->rmap->n;
627   PetscBool      isBAIJ,isSELL;
628 
629   PetscFunctionBegin;
630   /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian;
631    bcols is chosen s.t. dy-array takes 50% of memory space as mat */
632   ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
633   ierr = PetscObjectTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);CHKERRQ(ierr);
634   ierr = PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL);CHKERRQ(ierr);
635   if (isBAIJ || m == 0) {
636     c->brows = m;
637     c->bcols = 1;
638   } else if (isSELL) {
639     /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
640     Mat_MPISELL *sell=(Mat_MPISELL*)mat->data;
641     Mat_SeqSELL *spA,*spB;
642     Mat        A,B;
643     PetscInt   nz,brows,bcols;
644     PetscReal  mem;
645 
646     bs    = 1; /* only bs=1 is supported for MPISELL matrix */
647 
648     A = sell->A;  spA = (Mat_SeqSELL*)A->data;
649     B = sell->B;  spB = (Mat_SeqSELL*)B->data;
650     nz = spA->nz + spB->nz; /* total local nonzero entries of mat */
651     mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
652     bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
653     brows = 1000/bcols;
654     if (bcols > nis) bcols = nis;
655     if (brows == 0 || brows > m) brows = m;
656     c->brows = brows;
657     c->bcols = bcols;
658   } else { /* mpiaij matrix */
659     /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
660     Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
661     Mat_SeqAIJ *spA,*spB;
662     Mat        A,B;
663     PetscInt   nz,brows,bcols;
664     PetscReal  mem;
665 
666     bs    = 1; /* only bs=1 is supported for MPIAIJ matrix */
667 
668     A = aij->A;  spA = (Mat_SeqAIJ*)A->data;
669     B = aij->B;  spB = (Mat_SeqAIJ*)B->data;
670     nz = spA->nz + spB->nz; /* total local nonzero entries of mat */
671     mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
672     bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
673     brows = 1000/bcols;
674     if (bcols > nis) bcols = nis;
675     if (brows == 0 || brows > m) brows = m;
676     c->brows = brows;
677     c->bcols = bcols;
678   }
679 
680   c->M       = mat->rmap->N/bs;         /* set the global rows and columns and local rows */
681   c->N       = mat->cmap->N/bs;
682   c->m       = mat->rmap->n/bs;
683   c->rstart  = mat->rmap->rstart/bs;
684   c->ncolors = nis;
685   PetscFunctionReturn(0);
686 }
687