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
MatFDColoringApply_BAIJ(Mat J,MatFDColoring coloring,Vec x1,void * sctx)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 = VecBindToCPU(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 = VecBindToCPU(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 = VecBindToCPU(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 */
MatFDColoringApply_AIJ(Mat J,MatFDColoring coloring,Vec x1,void * sctx)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 = VecBindToCPU(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 #if defined(PETSC_HAVE_VIENNACL) || defined(PETSC_HAVE_CUDA)
356 if (J->offloadmask != PETSC_OFFLOAD_UNALLOCATED) J->offloadmask = PETSC_OFFLOAD_CPU;
357 #endif
358 ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
359 ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
360 if (vscale) {
361 ierr = VecRestoreArray(vscale,&vscale_array);CHKERRQ(ierr);
362 }
363 coloring->currentcolor = -1;
364 ierr = VecBindToCPU(x1,PETSC_FALSE);CHKERRQ(ierr);
365 PetscFunctionReturn(0);
366 }
367
MatFDColoringSetUp_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)368 PetscErrorCode MatFDColoringSetUp_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
369 {
370 PetscErrorCode ierr;
371 PetscMPIInt size,*ncolsonproc,*disp,nn;
372 PetscInt i,n,nrows,nrows_i,j,k,m,ncols,col,*rowhit,cstart,cend,colb;
373 const PetscInt *is,*A_ci,*A_cj,*B_ci,*B_cj,*row=NULL,*ltog=NULL;
374 PetscInt nis=iscoloring->n,nctot,*cols,tmp = 0;
375 ISLocalToGlobalMapping map=mat->cmap->mapping;
376 PetscInt ctype=c->ctype,*spidxA,*spidxB,nz,bs,bs2,spidx;
377 Mat A,B;
378 PetscScalar *A_val,*B_val,**valaddrhit;
379 MatEntry *Jentry;
380 MatEntry2 *Jentry2;
381 PetscBool isBAIJ,isSELL;
382 PetscInt bcols=c->bcols;
383 #if defined(PETSC_USE_CTABLE)
384 PetscTable colmap=NULL;
385 #else
386 PetscInt *colmap=NULL; /* local col number of off-diag col */
387 #endif
388
389 PetscFunctionBegin;
390 if (ctype == IS_COLORING_LOCAL) {
391 if (!map) SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_ARG_INCOMP,"When using ghosted differencing matrix must have local to global mapping provided with MatSetLocalToGlobalMapping");
392 ierr = ISLocalToGlobalMappingGetIndices(map,<og);CHKERRQ(ierr);
393 }
394
395 ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
396 ierr = PetscObjectBaseTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);CHKERRQ(ierr);
397 ierr = PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL);CHKERRQ(ierr);
398 if (isBAIJ) {
399 Mat_MPIBAIJ *baij=(Mat_MPIBAIJ*)mat->data;
400 Mat_SeqBAIJ *spA,*spB;
401 A = baij->A; spA = (Mat_SeqBAIJ*)A->data; A_val = spA->a;
402 B = baij->B; spB = (Mat_SeqBAIJ*)B->data; B_val = spB->a;
403 nz = spA->nz + spB->nz; /* total nonzero entries of mat */
404 if (!baij->colmap) {
405 ierr = MatCreateColmap_MPIBAIJ_Private(mat);CHKERRQ(ierr);
406 }
407 colmap = baij->colmap;
408 ierr = MatGetColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
409 ierr = MatGetColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
410
411 if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
412 PetscInt *garray;
413 ierr = PetscMalloc1(B->cmap->n,&garray);CHKERRQ(ierr);
414 for (i=0; i<baij->B->cmap->n/bs; i++) {
415 for (j=0; j<bs; j++) {
416 garray[i*bs+j] = bs*baij->garray[i]+j;
417 }
418 }
419 ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,garray,&c->vscale);CHKERRQ(ierr);
420 ierr = VecBindToCPU(c->vscale,PETSC_TRUE);CHKERRQ(ierr);
421 ierr = PetscFree(garray);CHKERRQ(ierr);
422 }
423 } else if (isSELL) {
424 Mat_MPISELL *sell=(Mat_MPISELL*)mat->data;
425 Mat_SeqSELL *spA,*spB;
426 A = sell->A; spA = (Mat_SeqSELL*)A->data; A_val = spA->val;
427 B = sell->B; spB = (Mat_SeqSELL*)B->data; B_val = spB->val;
428 nz = spA->nz + spB->nz; /* total nonzero entries of mat */
429 if (!sell->colmap) {
430 /* Allow access to data structures of local part of matrix
431 - creates aij->colmap which maps global column number to local number in part B */
432 ierr = MatCreateColmap_MPISELL_Private(mat);CHKERRQ(ierr);
433 }
434 colmap = sell->colmap;
435 ierr = MatGetColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
436 ierr = MatGetColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
437
438 bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */
439
440 if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
441 ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,sell->garray,&c->vscale);CHKERRQ(ierr);
442 ierr = VecBindToCPU(c->vscale,PETSC_TRUE);CHKERRQ(ierr);
443 }
444 } else {
445 Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
446 Mat_SeqAIJ *spA,*spB;
447 A = aij->A; spA = (Mat_SeqAIJ*)A->data; A_val = spA->a;
448 B = aij->B; spB = (Mat_SeqAIJ*)B->data; B_val = spB->a;
449 nz = spA->nz + spB->nz; /* total nonzero entries of mat */
450 if (!aij->colmap) {
451 /* Allow access to data structures of local part of matrix
452 - creates aij->colmap which maps global column number to local number in part B */
453 ierr = MatCreateColmap_MPIAIJ_Private(mat);CHKERRQ(ierr);
454 }
455 colmap = aij->colmap;
456 ierr = MatGetColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
457 ierr = MatGetColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
458
459 bs = 1; /* only bs=1 is supported for non MPIBAIJ matrix */
460
461 if (ctype == IS_COLORING_GLOBAL && c->htype[0] == 'd') { /* create vscale for storing dx */
462 ierr = VecCreateGhost(PetscObjectComm((PetscObject)mat),mat->cmap->n,PETSC_DETERMINE,B->cmap->n,aij->garray,&c->vscale);CHKERRQ(ierr);
463 ierr = VecBindToCPU(c->vscale,PETSC_TRUE);CHKERRQ(ierr);
464 }
465 }
466
467 m = mat->rmap->n/bs;
468 cstart = mat->cmap->rstart/bs;
469 cend = mat->cmap->rend/bs;
470
471 ierr = PetscMalloc2(nis,&c->ncolumns,nis,&c->columns);CHKERRQ(ierr);
472 ierr = PetscMalloc1(nis,&c->nrows);CHKERRQ(ierr);
473 ierr = PetscLogObjectMemory((PetscObject)c,3*nis*sizeof(PetscInt));CHKERRQ(ierr);
474
475 if (c->htype[0] == 'd') {
476 ierr = PetscMalloc1(nz,&Jentry);CHKERRQ(ierr);
477 ierr = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry));CHKERRQ(ierr);
478 c->matentry = Jentry;
479 } else if (c->htype[0] == 'w') {
480 ierr = PetscMalloc1(nz,&Jentry2);CHKERRQ(ierr);
481 ierr = PetscLogObjectMemory((PetscObject)c,nz*sizeof(MatEntry2));CHKERRQ(ierr);
482 c->matentry2 = Jentry2;
483 } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"htype is not supported");
484
485 ierr = PetscMalloc2(m+1,&rowhit,m+1,&valaddrhit);CHKERRQ(ierr);
486 nz = 0;
487 ierr = ISColoringGetIS(iscoloring,PETSC_OWN_POINTER, PETSC_IGNORE,&c->isa);CHKERRQ(ierr);
488
489 if (ctype == IS_COLORING_GLOBAL) {
490 ierr = MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size);CHKERRQ(ierr);
491 ierr = PetscMalloc2(size,&ncolsonproc,size,&disp);CHKERRQ(ierr);
492 }
493
494 for (i=0; i<nis; i++) { /* for each local color */
495 ierr = ISGetLocalSize(c->isa[i],&n);CHKERRQ(ierr);
496 ierr = ISGetIndices(c->isa[i],&is);CHKERRQ(ierr);
497
498 c->ncolumns[i] = n; /* local number of columns of this color on this process */
499 c->columns[i] = (PetscInt*)is;
500
501 if (ctype == IS_COLORING_GLOBAL) {
502 /* Determine nctot, the total (parallel) number of columns of this color */
503 /* ncolsonproc[j]: local ncolumns on proc[j] of this color */
504 ierr = PetscMPIIntCast(n,&nn);CHKERRQ(ierr);
505 ierr = MPI_Allgather(&nn,1,MPI_INT,ncolsonproc,1,MPI_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr);
506 nctot = 0; for (j=0; j<size; j++) nctot += ncolsonproc[j];
507 if (!nctot) {
508 ierr = PetscInfo(mat,"Coloring of matrix has some unneeded colors with no corresponding rows\n");CHKERRQ(ierr);
509 }
510
511 disp[0] = 0;
512 for (j=1; j<size; j++) {
513 disp[j] = disp[j-1] + ncolsonproc[j-1];
514 }
515
516 /* Get cols, the complete list of columns for this color on each process */
517 ierr = PetscMalloc1(nctot+1,&cols);CHKERRQ(ierr);
518 ierr = MPI_Allgatherv((void*)is,n,MPIU_INT,cols,ncolsonproc,disp,MPIU_INT,PetscObjectComm((PetscObject)mat));CHKERRQ(ierr);
519 } else if (ctype == IS_COLORING_LOCAL) {
520 /* Determine local number of columns of this color on this process, including ghost points */
521 nctot = n;
522 cols = (PetscInt*)is;
523 } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not provided for this MatFDColoring type");
524
525 /* Mark all rows affect by these columns */
526 ierr = PetscArrayzero(rowhit,m);CHKERRQ(ierr);
527 bs2 = bs*bs;
528 nrows_i = 0;
529 for (j=0; j<nctot; j++) { /* loop over columns*/
530 if (ctype == IS_COLORING_LOCAL) {
531 col = ltog[cols[j]];
532 } else {
533 col = cols[j];
534 }
535 if (col >= cstart && col < cend) { /* column is in A, diagonal block of mat */
536 tmp = A_ci[col-cstart];
537 row = A_cj + tmp;
538 nrows = A_ci[col-cstart+1] - tmp;
539 nrows_i += nrows;
540
541 /* loop over columns of A marking them in rowhit */
542 for (k=0; k<nrows; k++) {
543 /* set valaddrhit for part A */
544 spidx = bs2*spidxA[tmp + k];
545 valaddrhit[*row] = &A_val[spidx];
546 rowhit[*row++] = col - cstart + 1; /* local column index */
547 }
548 } else { /* column is in B, off-diagonal block of mat */
549 #if defined(PETSC_USE_CTABLE)
550 ierr = PetscTableFind(colmap,col+1,&colb);CHKERRQ(ierr);
551 colb--;
552 #else
553 colb = colmap[col] - 1; /* local column index */
554 #endif
555 if (colb == -1) {
556 nrows = 0;
557 } else {
558 colb = colb/bs;
559 tmp = B_ci[colb];
560 row = B_cj + tmp;
561 nrows = B_ci[colb+1] - tmp;
562 }
563 nrows_i += nrows;
564 /* loop over columns of B marking them in rowhit */
565 for (k=0; k<nrows; k++) {
566 /* set valaddrhit for part B */
567 spidx = bs2*spidxB[tmp + k];
568 valaddrhit[*row] = &B_val[spidx];
569 rowhit[*row++] = colb + 1 + cend - cstart; /* local column index */
570 }
571 }
572 }
573 c->nrows[i] = nrows_i;
574
575 if (c->htype[0] == 'd') {
576 for (j=0; j<m; j++) {
577 if (rowhit[j]) {
578 Jentry[nz].row = j; /* local row index */
579 Jentry[nz].col = rowhit[j] - 1; /* local column index */
580 Jentry[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */
581 nz++;
582 }
583 }
584 } else { /* c->htype == 'wp' */
585 for (j=0; j<m; j++) {
586 if (rowhit[j]) {
587 Jentry2[nz].row = j; /* local row index */
588 Jentry2[nz].valaddr = valaddrhit[j]; /* address of mat value for this entry */
589 nz++;
590 }
591 }
592 }
593 if (ctype == IS_COLORING_GLOBAL) {
594 ierr = PetscFree(cols);CHKERRQ(ierr);
595 }
596 }
597 if (ctype == IS_COLORING_GLOBAL) {
598 ierr = PetscFree2(ncolsonproc,disp);CHKERRQ(ierr);
599 }
600
601 if (bcols > 1) { /* reorder Jentry for faster MatFDColoringApply() */
602 ierr = MatFDColoringSetUpBlocked_AIJ_Private(mat,c,nz);CHKERRQ(ierr);
603 }
604
605 if (isBAIJ) {
606 ierr = MatRestoreColumnIJ_SeqBAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
607 ierr = MatRestoreColumnIJ_SeqBAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
608 ierr = PetscMalloc1(bs*mat->rmap->n,&c->dy);CHKERRQ(ierr);
609 } else if (isSELL) {
610 ierr = MatRestoreColumnIJ_SeqSELL_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
611 ierr = MatRestoreColumnIJ_SeqSELL_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
612 } else {
613 ierr = MatRestoreColumnIJ_SeqAIJ_Color(A,0,PETSC_FALSE,PETSC_FALSE,&ncols,&A_ci,&A_cj,&spidxA,NULL);CHKERRQ(ierr);
614 ierr = MatRestoreColumnIJ_SeqAIJ_Color(B,0,PETSC_FALSE,PETSC_FALSE,&ncols,&B_ci,&B_cj,&spidxB,NULL);CHKERRQ(ierr);
615 }
616
617 ierr = ISColoringRestoreIS(iscoloring,PETSC_OWN_POINTER,&c->isa);CHKERRQ(ierr);
618 ierr = PetscFree2(rowhit,valaddrhit);CHKERRQ(ierr);
619
620 if (ctype == IS_COLORING_LOCAL) {
621 ierr = ISLocalToGlobalMappingRestoreIndices(map,<og);CHKERRQ(ierr);
622 }
623 ierr = PetscInfo3(c,"ncolors %D, brows %D and bcols %D are used.\n",c->ncolors,c->brows,c->bcols);CHKERRQ(ierr);
624 PetscFunctionReturn(0);
625 }
626
MatFDColoringCreate_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)627 PetscErrorCode MatFDColoringCreate_MPIXAIJ(Mat mat,ISColoring iscoloring,MatFDColoring c)
628 {
629 PetscErrorCode ierr;
630 PetscInt bs,nis=iscoloring->n,m=mat->rmap->n;
631 PetscBool isBAIJ,isSELL;
632
633 PetscFunctionBegin;
634 /* set default brows and bcols for speedup inserting the dense matrix into sparse Jacobian;
635 bcols is chosen s.t. dy-array takes 50% of memory space as mat */
636 ierr = MatGetBlockSize(mat,&bs);CHKERRQ(ierr);
637 ierr = PetscObjectBaseTypeCompare((PetscObject)mat,MATMPIBAIJ,&isBAIJ);CHKERRQ(ierr);
638 ierr = PetscObjectTypeCompare((PetscObject)mat,MATMPISELL,&isSELL);CHKERRQ(ierr);
639 if (isBAIJ || m == 0) {
640 c->brows = m;
641 c->bcols = 1;
642 } else if (isSELL) {
643 /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
644 Mat_MPISELL *sell=(Mat_MPISELL*)mat->data;
645 Mat_SeqSELL *spA,*spB;
646 Mat A,B;
647 PetscInt nz,brows,bcols;
648 PetscReal mem;
649
650 bs = 1; /* only bs=1 is supported for MPISELL matrix */
651
652 A = sell->A; spA = (Mat_SeqSELL*)A->data;
653 B = sell->B; spB = (Mat_SeqSELL*)B->data;
654 nz = spA->nz + spB->nz; /* total local nonzero entries of mat */
655 mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
656 bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
657 brows = 1000/bcols;
658 if (bcols > nis) bcols = nis;
659 if (brows == 0 || brows > m) brows = m;
660 c->brows = brows;
661 c->bcols = bcols;
662 } else { /* mpiaij matrix */
663 /* bcols is chosen s.t. dy-array takes 50% of local memory space as mat */
664 Mat_MPIAIJ *aij=(Mat_MPIAIJ*)mat->data;
665 Mat_SeqAIJ *spA,*spB;
666 Mat A,B;
667 PetscInt nz,brows,bcols;
668 PetscReal mem;
669
670 bs = 1; /* only bs=1 is supported for MPIAIJ matrix */
671
672 A = aij->A; spA = (Mat_SeqAIJ*)A->data;
673 B = aij->B; spB = (Mat_SeqAIJ*)B->data;
674 nz = spA->nz + spB->nz; /* total local nonzero entries of mat */
675 mem = nz*(sizeof(PetscScalar) + sizeof(PetscInt)) + 3*m*sizeof(PetscInt);
676 bcols = (PetscInt)(0.5*mem /(m*sizeof(PetscScalar)));
677 brows = 1000/bcols;
678 if (bcols > nis) bcols = nis;
679 if (brows == 0 || brows > m) brows = m;
680 c->brows = brows;
681 c->bcols = bcols;
682 }
683
684 c->M = mat->rmap->N/bs; /* set the global rows and columns and local rows */
685 c->N = mat->cmap->N/bs;
686 c->m = mat->rmap->n/bs;
687 c->rstart = mat->rmap->rstart/bs;
688 c->ncolors = nis;
689 PetscFunctionReturn(0);
690 }
691
692 /*@C
693
694 MatFDColoringSetValues - takes a matrix in compressed color format and enters the matrix into a PETSc Mat
695
696 Collective on J
697
698 Input Parameters:
699 + J - the sparse matrix
700 . coloring - created with MatFDColoringCreate() and a local coloring
701 - y - column major storage of matrix values with one color of values per column, the number of rows of y should match
702 the number of local rows of J and the number of columns is the number of colors.
703
704 Level: intermediate
705
706 Notes: the matrix in compressed color format may come from an Automatic Differentiation code
707
708 The code will be slightly faster if MatFDColoringSetBlockSize(coloring,PETSC_DEFAULT,nc); is called immediately after creating the coloring
709
710 .seealso: MatFDColoringCreate(), ISColoring, ISColoringCreate(), ISColoringSetType(), IS_COLORING_LOCAL, MatFDColoringSetBlockSize()
711
712 @*/
MatFDColoringSetValues(Mat J,MatFDColoring coloring,const PetscScalar * y)713 PetscErrorCode MatFDColoringSetValues(Mat J,MatFDColoring coloring,const PetscScalar *y)
714 {
715 PetscErrorCode ierr;
716 MatEntry2 *Jentry2;
717 PetscInt row,i,nrows_k,l,ncolors,nz = 0,bcols,nbcols = 0;
718 const PetscInt *nrows;
719 PetscBool eq;
720
721 PetscFunctionBegin;
722 PetscValidHeaderSpecific(J,MAT_CLASSID,1);
723 PetscValidHeaderSpecific(coloring,MAT_FDCOLORING_CLASSID,2);
724 ierr = PetscObjectCompareId((PetscObject)J,coloring->matid,&eq);CHKERRQ(ierr);
725 if (!eq) SETERRQ(PetscObjectComm((PetscObject)J),PETSC_ERR_ARG_WRONG,"Matrix used with MatFDColoringSetValues() must be that used with MatFDColoringCreate()");
726 Jentry2 = coloring->matentry2;
727 nrows = coloring->nrows;
728 ncolors = coloring->ncolors;
729 bcols = coloring->bcols;
730
731 for (i=0; i<ncolors; i+=bcols) {
732 nrows_k = nrows[nbcols++];
733 for (l=0; l<nrows_k; l++) {
734 row = Jentry2[nz].row; /* local row index */
735 *(Jentry2[nz++].valaddr) = y[row];
736 }
737 y += bcols*coloring->m;
738 }
739 ierr = MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
740 ierr = MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
741 PetscFunctionReturn(0);
742 }
743