Actual source code: matpreallocator.c

petsc-main 2021-04-20
Report Typos and Errors
  1: #include <petsc/private/matimpl.h>
  2: #include <petsc/private/hashsetij.h>

  4: typedef struct {
  5:   PetscHSetIJ ht;
  6:   PetscInt   *dnz, *onz;
  7:   PetscInt   *dnzu, *onzu;
  8:   PetscBool   nooffproc;
  9: } Mat_Preallocator;

 11: PetscErrorCode MatDestroy_Preallocator(Mat A)
 12: {
 13:   Mat_Preallocator *p = (Mat_Preallocator *) A->data;
 14:   PetscErrorCode    ierr;

 17:   MatStashDestroy_Private(&A->stash);
 18:   PetscHSetIJDestroy(&p->ht);
 19:   PetscFree4(p->dnz, p->onz, p->dnzu, p->onzu);
 20:   PetscFree(A->data);
 21:   PetscObjectChangeTypeName((PetscObject) A, NULL);
 22:   PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", NULL);
 23:   return(0);
 24: }

 26: PetscErrorCode MatSetUp_Preallocator(Mat A)
 27: {
 28:   Mat_Preallocator *p = (Mat_Preallocator *) A->data;
 29:   PetscInt          m, bs, mbs;
 30:   PetscErrorCode    ierr;

 33:   PetscLayoutSetUp(A->rmap);
 34:   PetscLayoutSetUp(A->cmap);
 35:   MatGetLocalSize(A, &m, NULL);
 36:   PetscHSetIJCreate(&p->ht);
 37:   MatGetBlockSize(A, &bs);
 38:   /* Do not bother bstash since MatPreallocator does not implement MatSetValuesBlocked */
 39:   MatStashCreate_Private(PetscObjectComm((PetscObject) A), 1, &A->stash);
 40:   /* arrays are for blocked rows/cols */
 41:   mbs  = m/bs;
 42:   PetscCalloc4(mbs, &p->dnz, mbs, &p->onz, mbs, &p->dnzu, mbs, &p->onzu);
 43:   return(0);
 44: }

 46: PetscErrorCode MatSetValues_Preallocator(Mat A, PetscInt m, const PetscInt *rows, PetscInt n, const PetscInt *cols, const PetscScalar *values, InsertMode addv)
 47: {
 48:   Mat_Preallocator *p = (Mat_Preallocator *) A->data;
 49:   PetscInt          rStart, rEnd, r, cStart, cEnd, c, bs;
 50:   PetscErrorCode    ierr;

 53:   MatGetBlockSize(A, &bs);
 54:   MatGetOwnershipRange(A, &rStart, &rEnd);
 55:   MatGetOwnershipRangeColumn(A, &cStart, &cEnd);
 56:   for (r = 0; r < m; ++r) {
 57:     PetscHashIJKey key;
 58:     PetscBool      missing;

 60:     key.i = rows[r];
 61:     if (key.i < 0) continue;
 62:     if ((key.i < rStart) || (key.i >= rEnd)) {
 63:       MatStashValuesRow_Private(&A->stash, key.i, n, cols, values, PETSC_FALSE);
 64:     } else { /* Hash table is for blocked rows/cols */
 65:       key.i = rows[r]/bs;
 66:       for (c = 0; c < n; ++c) {
 67:         key.j = cols[c]/bs;
 68:         if (key.j < 0) continue;
 69:         PetscHSetIJQueryAdd(p->ht, key, &missing);
 70:         if (missing) {
 71:           if ((key.j >= cStart/bs) && (key.j < cEnd/bs)) {
 72:             ++p->dnz[key.i-rStart/bs];
 73:             if (key.j >= key.i) ++p->dnzu[key.i-rStart/bs];
 74:           } else {
 75:             ++p->onz[key.i-rStart/bs];
 76:             if (key.j >= key.i) ++p->onzu[key.i-rStart/bs];
 77:           }
 78:         }
 79:       }
 80:     }
 81:   }
 82:   return(0);
 83: }

 85: PetscErrorCode MatAssemblyBegin_Preallocator(Mat A, MatAssemblyType type)
 86: {
 87:   PetscInt       nstash, reallocs;

 91:   MatStashScatterBegin_Private(A, &A->stash, A->rmap->range);
 92:   MatStashGetInfo_Private(&A->stash, &nstash, &reallocs);
 93:   PetscInfo2(A, "Stash has %D entries, uses %D mallocs.\n", nstash, reallocs);
 94:   return(0);
 95: }

 97: PetscErrorCode MatAssemblyEnd_Preallocator(Mat A, MatAssemblyType type)
 98: {
 99:   PetscScalar      *val;
100:   PetscInt         *row, *col;
101:   PetscInt         i, j, rstart, ncols, flg;
102:   PetscMPIInt      n;
103:   Mat_Preallocator *p = (Mat_Preallocator *) A->data;
104:   PetscErrorCode   ierr;

107:   p->nooffproc = PETSC_TRUE;
108:   while (1) {
109:     MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg);
110:     if (flg) p->nooffproc = PETSC_FALSE;
111:     if (!flg) break;

113:     for (i = 0; i < n;) {
114:       /* Now identify the consecutive vals belonging to the same row */
115:       for (j = i, rstart = row[j]; j < n; j++) {
116:         if (row[j] != rstart) break;
117:       }
118:       if (j < n) ncols = j-i;
119:       else       ncols = n-i;
120:       /* Now assemble all these values with a single function call */
121:       MatSetValues_Preallocator(A, 1, row+i, ncols, col+i, val+i, INSERT_VALUES);
122:       i = j;
123:     }
124:   }
125:   MatStashScatterEnd_Private(&A->stash);
126:   MPIU_Allreduce(MPI_IN_PLACE,&p->nooffproc,1,MPIU_BOOL,MPI_LAND,PetscObjectComm((PetscObject)A));
127:   return(0);
128: }

130: PetscErrorCode MatView_Preallocator(Mat A, PetscViewer viewer)
131: {
133:   return(0);
134: }

136: PetscErrorCode MatSetOption_Preallocator(Mat A, MatOption op, PetscBool flg)
137: {
139:   return(0);
140: }

142: PetscErrorCode MatPreallocatorPreallocate_Preallocator(Mat mat, PetscBool fill, Mat A)
143: {
144:   Mat_Preallocator *p = (Mat_Preallocator *) mat->data;
145:   PetscInt          bs;
146:   PetscErrorCode    ierr;

149:   MatGetBlockSize(mat, &bs);
150:   MatXAIJSetPreallocation(A, bs, p->dnz, p->onz, p->dnzu, p->onzu);
151:   MatSetUp(A);
152:   MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);
153:   MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, p->nooffproc);
154:   if (fill) {
155:     PetscHashIter  hi;
156:     PetscHashIJKey key;
157:     PetscScalar    *zeros;

159:     PetscCalloc1(bs*bs,&zeros);

161:     PetscHashIterBegin(p->ht,hi);
162:     while (!PetscHashIterAtEnd(p->ht,hi)) {
163:       PetscHashIterGetKey(p->ht,hi,key);
164:       PetscHashIterNext(p->ht,hi);
165:       MatSetValuesBlocked(A,1,&key.i,1,&key.j,zeros,INSERT_VALUES);
166:     }
167:     PetscFree(zeros);

169:     MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
170:     MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
171:   }
172:   return(0);
173: }

175: /*@
176:   MatPreallocatorPreallocate - Preallocates the A matrix, using information from mat, optionally filling A with zeros

178:   Input Parameter:
179: + mat  - the preallocator
180: . fill - fill the matrix with zeros
181: - A    - the matrix to be preallocated

183:   Notes:
184:   This Mat implementation provides a helper utility to define the correct
185:   preallocation data for a given nonzero structure. Use this object like a
186:   regular matrix, e.g. loop over the nonzero structure of the matrix and
187:   call MatSetValues() or MatSetValuesBlocked() to indicate the nonzero locations.
188:   The matrix entires provided to MatSetValues() will be ignored, it only uses
189:   the row / col indices provided to determine the information required to be
190:   passed to MatXAIJSetPreallocation(). Once you have looped over the nonzero
191:   structure, you must call MatAssemblyBegin(), MatAssemblyEnd() on mat.

193:   After you have assembled the preallocator matrix (mat), call MatPreallocatorPreallocate()
194:   to define the preallocation information on the matrix (A). Setting the parameter
195:   fill = PETSC_TRUE will insert zeros into the matrix A. Internally MatPreallocatorPreallocate()
196:   will call MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);

198:   Level: advanced

200: .seealso: MATPREALLOCATOR
201: @*/
202: PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A)
203: {

209:   PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat,PetscBool,Mat),(mat,fill,A));
210:   return(0);
211: }

213: /*MC
214:    MATPREALLOCATOR - MATPREALLOCATOR = "preallocator" - A matrix type to be used for computing a matrix preallocation.

216:    Operations Provided:
217: .  MatSetValues()

219:    Options Database Keys:
220: . -mat_type preallocator - sets the matrix type to "preallocator" during a call to MatSetFromOptions()

222:   Level: advanced

224: .seealso: Mat

226: M*/

228: PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A)
229: {
230:   Mat_Preallocator *p;
231:   PetscErrorCode    ierr;

234:   PetscNewLog(A, &p);
235:   A->data = (void *) p;

237:   p->ht   = NULL;
238:   p->dnz  = NULL;
239:   p->onz  = NULL;
240:   p->dnzu = NULL;
241:   p->onzu = NULL;

243:   /* matrix ops */
244:   PetscMemzero(A->ops, sizeof(struct _MatOps));

246:   A->ops->destroy       = MatDestroy_Preallocator;
247:   A->ops->setup         = MatSetUp_Preallocator;
248:   A->ops->setvalues     = MatSetValues_Preallocator;
249:   A->ops->assemblybegin = MatAssemblyBegin_Preallocator;
250:   A->ops->assemblyend   = MatAssemblyEnd_Preallocator;
251:   A->ops->view          = MatView_Preallocator;
252:   A->ops->setoption     = MatSetOption_Preallocator;
253:   A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */

255:   /* special MATPREALLOCATOR functions */
256:   PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator);
257:   PetscObjectChangeTypeName((PetscObject) A, MATPREALLOCATOR);
258:   return(0);
259: }