xref: /petsc/src/mat/utils/axpy.c (revision 0e9ea286fb2a802ae93ee195655c93a598b51114)
1 
2 #include <petsc/private/matimpl.h>  /*I   "petscmat.h"  I*/
3 
4 /*@
5    MatAXPY - Computes Y = a*X + Y.
6 
7    Logically  Collective on Mat
8 
9    Input Parameters:
10 +  a - the scalar multiplier
11 .  X - the first matrix
12 .  Y - the second matrix
13 -  str - either SAME_NONZERO_PATTERN, DIFFERENT_NONZERO_PATTERN
14          or SUBSET_NONZERO_PATTERN (nonzeros of X is a subset of Y's)
15 
16    Level: intermediate
17 
18 .keywords: matrix, add
19 
20 .seealso: MatAYPX()
21  @*/
22 PetscErrorCode MatAXPY(Mat Y,PetscScalar a,Mat X,MatStructure str)
23 {
24   PetscErrorCode ierr;
25   PetscInt       M1,M2,N1,N2;
26   PetscInt       m1,m2,n1,n2;
27   PetscBool      sametype;
28 
29   PetscFunctionBegin;
30   PetscValidHeaderSpecific(Y,MAT_CLASSID,1);
31   PetscValidLogicalCollectiveScalar(Y,a,2);
32   PetscValidHeaderSpecific(X,MAT_CLASSID,3);
33   PetscCheckSameComm(Y,1,X,3);
34   ierr = MatGetSize(X,&M1,&N1);CHKERRQ(ierr);
35   ierr = MatGetSize(Y,&M2,&N2);CHKERRQ(ierr);
36   ierr = MatGetLocalSize(X,&m1,&n1);CHKERRQ(ierr);
37   ierr = MatGetLocalSize(Y,&m2,&n2);CHKERRQ(ierr);
38   if (M1 != M2 || N1 != N2) SETERRQ4(PetscObjectComm((PetscObject)Y),PETSC_ERR_ARG_SIZ,"Non conforming matrix add: global sizes %D x %D, %D x %D",M1,M2,N1,N2);
39   if (m1 != m2 || n1 != n2) SETERRQ4(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Non conforming matrix add: local sizes %D x %D, %D x %D",m1,m2,n1,n2);
40 
41   ierr = PetscStrcmp(((PetscObject)X)->type_name,((PetscObject)Y)->type_name,&sametype);CHKERRQ(ierr);
42   ierr = PetscLogEventBegin(MAT_AXPY,Y,0,0,0);CHKERRQ(ierr);
43   if (Y->ops->axpy && sametype) {
44     ierr = (*Y->ops->axpy)(Y,a,X,str);CHKERRQ(ierr);
45   } else {
46     if (str != DIFFERENT_NONZERO_PATTERN) {
47       ierr = MatAXPY_Basic(Y,a,X,str);CHKERRQ(ierr);
48     } else {
49       Mat B;
50 
51       ierr = MatAXPY_Basic_Preallocate(Y,X,&B);CHKERRQ(ierr);
52       ierr = MatAXPY_BasicWithPreallocation(B,Y,a,X,str);CHKERRQ(ierr);
53       ierr = MatHeaderReplace(Y,&B);CHKERRQ(ierr);
54     }
55   }
56   ierr = PetscLogEventEnd(MAT_AXPY,Y,0,0,0);CHKERRQ(ierr);
57 #if defined(PETSC_HAVE_VIENNACL) || defined(PETSC_HAVE_CUDA)
58   if (Y->valid_GPU_matrix != PETSC_OFFLOAD_UNALLOCATED) {
59     Y->valid_GPU_matrix = PETSC_OFFLOAD_CPU;
60   }
61 #endif
62   PetscFunctionReturn(0);
63 }
64 
65 PetscErrorCode MatAXPY_Basic_Preallocate(Mat Y, Mat X, Mat *B)
66 {
67   PetscErrorCode ierr;
68 
69   PetscFunctionBegin;
70   /* look for any available faster alternative to the general preallocator */
71   PetscErrorCode (*preall)(Mat,Mat,Mat*) = NULL;
72   ierr = PetscObjectQueryFunction((PetscObject)Y,"MatAXPYGetPreallocation_C",&preall);CHKERRQ(ierr);
73   if (preall) {
74     ierr = (*preall)(Y,X,B);CHKERRQ(ierr);
75   } else { /* Use MatPrellocator, assumes same row-col distribution */
76     Mat      preallocator;
77     PetscInt r,rstart,rend;
78     PetscInt m,n,M,N;
79 
80     ierr = MatGetSize(Y,&M,&N);CHKERRQ(ierr);
81     ierr = MatGetLocalSize(Y,&m,&n);CHKERRQ(ierr);
82     ierr = MatCreate(PetscObjectComm((PetscObject)Y),&preallocator);CHKERRQ(ierr);
83     ierr = MatSetType(preallocator,MATPREALLOCATOR);CHKERRQ(ierr);
84     ierr = MatSetSizes(preallocator,m,n,M,N);CHKERRQ(ierr);
85     ierr = MatSetUp(preallocator);CHKERRQ(ierr);
86     ierr = MatGetOwnershipRange(preallocator,&rstart,&rend);CHKERRQ(ierr);
87     for (r = rstart; r < rend; ++r) {
88       PetscInt          ncols;
89       const PetscInt    *row;
90       const PetscScalar *vals;
91 
92       ierr = MatGetRow(Y,r,&ncols,&row,&vals);CHKERRQ(ierr);
93       ierr = MatSetValues(preallocator,1,&r,ncols,row,vals,INSERT_VALUES);CHKERRQ(ierr);
94       ierr = MatRestoreRow(Y,r,&ncols,&row,&vals);CHKERRQ(ierr);
95       ierr = MatGetRow(X,r,&ncols,&row,&vals);CHKERRQ(ierr);
96       ierr = MatSetValues(preallocator,1,&r,ncols,row,vals,INSERT_VALUES);CHKERRQ(ierr);
97       ierr = MatRestoreRow(X,r,&ncols,&row,&vals);CHKERRQ(ierr);
98     }
99     ierr = MatAssemblyBegin(preallocator,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
100     ierr = MatAssemblyEnd(preallocator,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
101 
102     ierr = MatCreate(PetscObjectComm((PetscObject)Y),B);CHKERRQ(ierr);
103     ierr = MatSetType(*B,((PetscObject)Y)->type_name);CHKERRQ(ierr);
104     ierr = MatSetSizes(*B,m,n,M,N);CHKERRQ(ierr);
105     ierr = MatPreallocatorPreallocate(preallocator,PETSC_FALSE,*B);CHKERRQ(ierr);
106     ierr = MatDestroy(&preallocator);CHKERRQ(ierr);
107   }
108   PetscFunctionReturn(0);
109 }
110 
111 PetscErrorCode MatAXPY_Basic(Mat Y,PetscScalar a,Mat X,MatStructure str)
112 {
113   PetscInt          i,start,end,j,ncols,m,n;
114   PetscErrorCode    ierr;
115   const PetscInt    *row;
116   PetscScalar       *val;
117   const PetscScalar *vals;
118 
119   PetscFunctionBegin;
120   ierr = MatGetSize(X,&m,&n);CHKERRQ(ierr);
121   ierr = MatGetOwnershipRange(X,&start,&end);CHKERRQ(ierr);
122   if (a == 1.0) {
123     for (i = start; i < end; i++) {
124       ierr = MatGetRow(X,i,&ncols,&row,&vals);CHKERRQ(ierr);
125       ierr = MatSetValues(Y,1,&i,ncols,row,vals,ADD_VALUES);CHKERRQ(ierr);
126       ierr = MatRestoreRow(X,i,&ncols,&row,&vals);CHKERRQ(ierr);
127     }
128   } else {
129     ierr = PetscMalloc1(n+1,&val);CHKERRQ(ierr);
130     for (i=start; i<end; i++) {
131       ierr = MatGetRow(X,i,&ncols,&row,&vals);CHKERRQ(ierr);
132       for (j=0; j<ncols; j++) {
133         val[j] = a*vals[j];
134       }
135       ierr = MatSetValues(Y,1,&i,ncols,row,val,ADD_VALUES);CHKERRQ(ierr);
136       ierr = MatRestoreRow(X,i,&ncols,&row,&vals);CHKERRQ(ierr);
137     }
138     ierr = PetscFree(val);CHKERRQ(ierr);
139   }
140   ierr = MatAssemblyBegin(Y,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
141   ierr = MatAssemblyEnd(Y,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
142   PetscFunctionReturn(0);
143 }
144 
145 PetscErrorCode MatAXPY_BasicWithPreallocation(Mat B,Mat Y,PetscScalar a,Mat X,MatStructure str)
146 {
147   PetscInt          i,start,end,j,ncols,m,n;
148   PetscErrorCode    ierr;
149   const PetscInt    *row;
150   PetscScalar       *val;
151   const PetscScalar *vals;
152 
153   PetscFunctionBegin;
154   ierr = MatGetSize(X,&m,&n);CHKERRQ(ierr);
155   ierr = MatGetOwnershipRange(X,&start,&end);CHKERRQ(ierr);
156   if (a == 1.0) {
157     for (i = start; i < end; i++) {
158       ierr = MatGetRow(Y,i,&ncols,&row,&vals);CHKERRQ(ierr);
159       ierr = MatSetValues(B,1,&i,ncols,row,vals,ADD_VALUES);CHKERRQ(ierr);
160       ierr = MatRestoreRow(Y,i,&ncols,&row,&vals);CHKERRQ(ierr);
161 
162       ierr = MatGetRow(X,i,&ncols,&row,&vals);CHKERRQ(ierr);
163       ierr = MatSetValues(B,1,&i,ncols,row,vals,ADD_VALUES);CHKERRQ(ierr);
164       ierr = MatRestoreRow(X,i,&ncols,&row,&vals);CHKERRQ(ierr);
165     }
166   } else {
167     ierr = PetscMalloc1(n+1,&val);CHKERRQ(ierr);
168     for (i=start; i<end; i++) {
169       ierr = MatGetRow(Y,i,&ncols,&row,&vals);CHKERRQ(ierr);
170       ierr = MatSetValues(B,1,&i,ncols,row,vals,ADD_VALUES);CHKERRQ(ierr);
171       ierr = MatRestoreRow(Y,i,&ncols,&row,&vals);CHKERRQ(ierr);
172 
173       ierr = MatGetRow(X,i,&ncols,&row,&vals);CHKERRQ(ierr);
174       for (j=0; j<ncols; j++) {
175         val[j] = a*vals[j];
176       }
177       ierr = MatSetValues(B,1,&i,ncols,row,val,ADD_VALUES);CHKERRQ(ierr);
178       ierr = MatRestoreRow(X,i,&ncols,&row,&vals);CHKERRQ(ierr);
179     }
180     ierr = PetscFree(val);CHKERRQ(ierr);
181   }
182   ierr = MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
183   ierr = MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
184   PetscFunctionReturn(0);
185 }
186 
187 /*@
188    MatShift - Computes Y =  Y + a I, where a is a PetscScalar and I is the identity matrix.
189 
190    Neighbor-wise Collective on Mat
191 
192    Input Parameters:
193 +  Y - the matrices
194 -  a - the PetscScalar
195 
196    Level: intermediate
197 
198    Notes:
199     If the matrix Y is missing some diagonal entries this routine can be very slow. To make it fast one should initially
200    fill the matrix so that all diagonal entries have a value (with a value of zero for those locations that would not have an
201    entry).
202 
203    To form Y = Y + diag(V) use MatDiagonalSet()
204 
205    Developers Note: If the local "diagonal part" of the matrix Y has no entries then the local diagonal part is
206     preallocated with 1 nonzero per row for the to be added values. This allows for fast shifting of an empty matrix.
207 
208 .keywords: matrix, add, shift
209 
210 .seealso: MatDiagonalSet(), MatScale(), MatDiagonalScale()
211  @*/
212 PetscErrorCode  MatShift(Mat Y,PetscScalar a)
213 {
214   PetscErrorCode ierr;
215 
216   PetscFunctionBegin;
217   PetscValidHeaderSpecific(Y,MAT_CLASSID,1);
218   if (!Y->assembled) SETERRQ(PetscObjectComm((PetscObject)Y),PETSC_ERR_ARG_WRONGSTATE,"Not for unassembled matrix");
219   if (Y->factortype) SETERRQ(PetscObjectComm((PetscObject)Y),PETSC_ERR_ARG_WRONGSTATE,"Not for factored matrix");
220   MatCheckPreallocated(Y,1);
221 
222   if (Y->ops->shift) {
223     ierr = (*Y->ops->shift)(Y,a);CHKERRQ(ierr);
224   } else {
225     ierr = MatShift_Basic(Y,a);CHKERRQ(ierr);
226   }
227 
228   ierr = PetscObjectStateIncrease((PetscObject)Y);CHKERRQ(ierr);
229 #if defined(PETSC_HAVE_VIENNACL) || defined(PETSC_HAVE_CUDA)
230   if (Y->valid_GPU_matrix != PETSC_OFFLOAD_UNALLOCATED) {
231     Y->valid_GPU_matrix = PETSC_OFFLOAD_CPU;
232   }
233 #endif
234   PetscFunctionReturn(0);
235 }
236 
237 PetscErrorCode  MatDiagonalSet_Default(Mat Y,Vec D,InsertMode is)
238 {
239   PetscErrorCode ierr;
240   PetscInt       i,start,end;
241   PetscScalar    *v;
242 
243   PetscFunctionBegin;
244   ierr = MatGetOwnershipRange(Y,&start,&end);CHKERRQ(ierr);
245   ierr = VecGetArray(D,&v);CHKERRQ(ierr);
246   for (i=start; i<end; i++) {
247     ierr = MatSetValues(Y,1,&i,1,&i,v+i-start,is);CHKERRQ(ierr);
248   }
249   ierr = VecRestoreArray(D,&v);CHKERRQ(ierr);
250   ierr = MatAssemblyBegin(Y,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
251   ierr = MatAssemblyEnd(Y,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
252   PetscFunctionReturn(0);
253 }
254 
255 /*@
256    MatDiagonalSet - Computes Y = Y + D, where D is a diagonal matrix
257    that is represented as a vector. Or Y[i,i] = D[i] if InsertMode is
258    INSERT_VALUES.
259 
260    Input Parameters:
261 +  Y - the input matrix
262 .  D - the diagonal matrix, represented as a vector
263 -  i - INSERT_VALUES or ADD_VALUES
264 
265    Neighbor-wise Collective on Mat and Vec
266 
267    Notes:
268     If the matrix Y is missing some diagonal entries this routine can be very slow. To make it fast one should initially
269    fill the matrix so that all diagonal entries have a value (with a value of zero for those locations that would not have an
270    entry).
271 
272    Level: intermediate
273 
274 .keywords: matrix, add, shift, diagonal
275 
276 .seealso: MatShift(), MatScale(), MatDiagonalScale()
277 @*/
278 PetscErrorCode  MatDiagonalSet(Mat Y,Vec D,InsertMode is)
279 {
280   PetscErrorCode ierr;
281   PetscInt       matlocal,veclocal;
282 
283   PetscFunctionBegin;
284   PetscValidHeaderSpecific(Y,MAT_CLASSID,1);
285   PetscValidHeaderSpecific(D,VEC_CLASSID,2);
286   ierr = MatGetLocalSize(Y,&matlocal,NULL);CHKERRQ(ierr);
287   ierr = VecGetLocalSize(D,&veclocal);CHKERRQ(ierr);
288   if (matlocal != veclocal) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Number local rows of matrix %D does not match that of vector for diagonal %D",matlocal,veclocal);
289   if (Y->ops->diagonalset) {
290     ierr = (*Y->ops->diagonalset)(Y,D,is);CHKERRQ(ierr);
291   } else {
292     ierr = MatDiagonalSet_Default(Y,D,is);CHKERRQ(ierr);
293   }
294   ierr = PetscObjectStateIncrease((PetscObject)Y);CHKERRQ(ierr);
295   PetscFunctionReturn(0);
296 }
297 
298 /*@
299    MatAYPX - Computes Y = a*Y + X.
300 
301    Logically on Mat
302 
303    Input Parameters:
304 +  a - the PetscScalar multiplier
305 .  Y - the first matrix
306 .  X - the second matrix
307 -  str - either SAME_NONZERO_PATTERN, DIFFERENT_NONZERO_PATTERN or SUBSET_NONZERO_PATTERN
308 
309    Level: intermediate
310 
311 .keywords: matrix, add
312 
313 .seealso: MatAXPY()
314  @*/
315 PetscErrorCode  MatAYPX(Mat Y,PetscScalar a,Mat X,MatStructure str)
316 {
317   PetscScalar    one = 1.0;
318   PetscErrorCode ierr;
319   PetscInt       mX,mY,nX,nY;
320 
321   PetscFunctionBegin;
322   PetscValidHeaderSpecific(X,MAT_CLASSID,3);
323   PetscValidHeaderSpecific(Y,MAT_CLASSID,1);
324   PetscValidLogicalCollectiveScalar(Y,a,2);
325   ierr = MatGetSize(X,&mX,&nX);CHKERRQ(ierr);
326   ierr = MatGetSize(X,&mY,&nY);CHKERRQ(ierr);
327   if (mX != mY || nX != nY) SETERRQ4(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Non conforming matrices: %D %D first %D %D second",mX,mY,nX,nY);
328 
329   ierr = MatScale(Y,a);CHKERRQ(ierr);
330   ierr = MatAXPY(Y,one,X,str);CHKERRQ(ierr);
331   PetscFunctionReturn(0);
332 }
333 
334 /*@
335     MatComputeExplicitOperator - Computes the explicit matrix
336 
337     Collective on Mat
338 
339     Input Parameter:
340 .   inmat - the matrix
341 
342     Output Parameter:
343 .   mat - the explict  operator
344 
345     Notes:
346     This computation is done by applying the operators to columns of the
347     identity matrix.
348 
349     Currently, this routine uses a dense matrix format when 1 processor
350     is used and a sparse format otherwise.  This routine is costly in general,
351     and is recommended for use only with relatively small systems.
352 
353     Level: advanced
354 
355 .keywords: Mat, compute, explicit, operator
356 @*/
357 PetscErrorCode  MatComputeExplicitOperator(Mat inmat,Mat *mat)
358 {
359   PetscErrorCode ierr;
360   MPI_Comm       comm;
361   PetscMPIInt    size;
362 
363   PetscFunctionBegin;
364   PetscValidHeaderSpecific(inmat,MAT_CLASSID,1);
365   PetscValidPointer(mat,2);
366 
367   ierr = PetscObjectGetComm((PetscObject)inmat,&comm);CHKERRQ(ierr);
368   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
369   ierr = MatConvert_Shell(inmat,size == 1 ? MATSEQDENSE : MATAIJ,MAT_INITIAL_MATRIX,mat);CHKERRQ(ierr);
370   PetscFunctionReturn(0);
371 }
372 
373 /*@
374     MatComputeExplicitOperatorTranspose - Computes the explicit matrix representation of
375         a give matrix that can apply MatMultTranspose()
376 
377     Collective on Mat
378 
379     Input Parameter:
380 .   inmat - the matrix
381 
382     Output Parameter:
383 .   mat - the explict  operator transposed
384 
385     Notes:
386     This computation is done by applying the transpose of the operator to columns of the
387     identity matrix.
388 
389     Currently, this routine uses a dense matrix format when 1 processor
390     is used and a sparse format otherwise.  This routine is costly in general,
391     and is recommended for use only with relatively small systems.
392 
393     Level: advanced
394 
395 .keywords: Mat, compute, explicit, operator
396 @*/
397 PetscErrorCode  MatComputeExplicitOperatorTranspose(Mat inmat,Mat *mat)
398 {
399   Vec            in,out;
400   PetscErrorCode ierr;
401   PetscInt       i,m,n,M,N,*rows,start,end;
402   MPI_Comm       comm;
403   PetscScalar    *array,zero = 0.0,one = 1.0;
404   PetscMPIInt    size;
405 
406   PetscFunctionBegin;
407   PetscValidHeaderSpecific(inmat,MAT_CLASSID,1);
408   PetscValidPointer(mat,2);
409 
410   ierr = PetscObjectGetComm((PetscObject)inmat,&comm);CHKERRQ(ierr);
411   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
412 
413   ierr = MatGetLocalSize(inmat,&m,&n);CHKERRQ(ierr);
414   ierr = MatGetSize(inmat,&M,&N);CHKERRQ(ierr);
415   ierr = MatCreateVecs(inmat,&in,&out);CHKERRQ(ierr);
416   ierr = VecSetOption(in,VEC_IGNORE_OFF_PROC_ENTRIES,PETSC_TRUE);CHKERRQ(ierr);
417   ierr = VecGetOwnershipRange(out,&start,&end);CHKERRQ(ierr);
418   ierr = PetscMalloc1(m,&rows);CHKERRQ(ierr);
419   for (i=0; i<m; i++) rows[i] = start + i;
420 
421   ierr = MatCreate(comm,mat);CHKERRQ(ierr);
422   ierr = MatSetSizes(*mat,m,n,M,N);CHKERRQ(ierr);
423   if (size == 1) {
424     ierr = MatSetType(*mat,MATSEQDENSE);CHKERRQ(ierr);
425     ierr = MatSeqDenseSetPreallocation(*mat,NULL);CHKERRQ(ierr);
426   } else {
427     ierr = MatSetType(*mat,MATMPIAIJ);CHKERRQ(ierr);
428     ierr = MatMPIAIJSetPreallocation(*mat,n,NULL,N-n,NULL);CHKERRQ(ierr);
429   }
430 
431   for (i=0; i<N; i++) {
432 
433     ierr = VecSet(in,zero);CHKERRQ(ierr);
434     ierr = VecSetValues(in,1,&i,&one,INSERT_VALUES);CHKERRQ(ierr);
435     ierr = VecAssemblyBegin(in);CHKERRQ(ierr);
436     ierr = VecAssemblyEnd(in);CHKERRQ(ierr);
437 
438     ierr = MatMultTranspose(inmat,in,out);CHKERRQ(ierr);
439 
440     ierr = VecGetArray(out,&array);CHKERRQ(ierr);
441     ierr = MatSetValues(*mat,m,rows,1,&i,array,INSERT_VALUES);CHKERRQ(ierr);
442     ierr = VecRestoreArray(out,&array);CHKERRQ(ierr);
443 
444   }
445   ierr = PetscFree(rows);CHKERRQ(ierr);
446   ierr = VecDestroy(&out);CHKERRQ(ierr);
447   ierr = VecDestroy(&in);CHKERRQ(ierr);
448   ierr = MatAssemblyBegin(*mat,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
449   ierr = MatAssemblyEnd(*mat,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
450   PetscFunctionReturn(0);
451 }
452 
453 /*@
454   MatChop - Set all values in the matrix less than the tolerance to zero
455 
456   Input Parameters:
457 + A   - The matrix
458 - tol - The zero tolerance
459 
460   Output Parameters:
461 . A - The chopped matrix
462 
463   Level: intermediate
464 
465 .seealso: MatCreate(), MatZeroEntries()
466  @*/
467 PetscErrorCode MatChop(Mat A, PetscReal tol)
468 {
469   PetscScalar    *newVals;
470   PetscInt       *newCols;
471   PetscInt       rStart, rEnd, numRows, maxRows, r, colMax = 0;
472   PetscErrorCode ierr;
473 
474   PetscFunctionBegin;
475   ierr = MatGetOwnershipRange(A, &rStart, &rEnd);CHKERRQ(ierr);
476   for (r = rStart; r < rEnd; ++r) {
477     PetscInt ncols;
478 
479     ierr   = MatGetRow(A, r, &ncols, NULL, NULL);CHKERRQ(ierr);
480     colMax = PetscMax(colMax, ncols);CHKERRQ(ierr);
481     ierr   = MatRestoreRow(A, r, &ncols, NULL, NULL);CHKERRQ(ierr);
482   }
483   numRows = rEnd - rStart;
484   ierr    = MPIU_Allreduce(&numRows, &maxRows, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)A));CHKERRQ(ierr);
485   ierr    = PetscMalloc2(colMax,&newCols,colMax,&newVals);CHKERRQ(ierr);
486   for (r = rStart; r < rStart+maxRows; ++r) {
487     const PetscScalar *vals;
488     const PetscInt    *cols;
489     PetscInt           ncols, newcols, c;
490 
491     if (r < rEnd) {
492       ierr = MatGetRow(A, r, &ncols, &cols, &vals);CHKERRQ(ierr);
493       for (c = 0; c < ncols; ++c) {
494         newCols[c] = cols[c];
495         newVals[c] = PetscAbsScalar(vals[c]) < tol ? 0.0 : vals[c];
496       }
497       newcols = ncols;
498       ierr = MatRestoreRow(A, r, &ncols, &cols, &vals);CHKERRQ(ierr);
499       ierr = MatSetValues(A, 1, &r, newcols, newCols, newVals, INSERT_VALUES);CHKERRQ(ierr);
500     }
501     ierr = MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
502     ierr = MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
503   }
504   ierr = PetscFree2(newCols,newVals);CHKERRQ(ierr);
505   PetscFunctionReturn(0);
506 }
507