Actual source code: matpreallocator.c

  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:   PetscBool   used;
 10: } Mat_Preallocator;

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

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

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

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

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

 48:   MatGetBlockSize(A, &bs);
 49:   MatGetOwnershipRange(A, &rStart, &rEnd);
 50:   MatGetOwnershipRangeColumn(A, &cStart, &cEnd);
 51:   for (r = 0; r < m; ++r) {
 52:     PetscHashIJKey key;
 53:     PetscBool      missing;

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

 80: PetscErrorCode MatAssemblyBegin_Preallocator(Mat A, MatAssemblyType type)
 81: {
 82:   PetscInt       nstash, reallocs;

 84:   MatStashScatterBegin_Private(A, &A->stash, A->rmap->range);
 85:   MatStashGetInfo_Private(&A->stash, &nstash, &reallocs);
 86:   PetscInfo(A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs);
 87:   return 0;
 88: }

 90: PetscErrorCode MatAssemblyEnd_Preallocator(Mat A, MatAssemblyType type)
 91: {
 92:   PetscScalar      *val;
 93:   PetscInt         *row, *col;
 94:   PetscInt         i, j, rstart, ncols, flg;
 95:   PetscMPIInt      n;
 96:   Mat_Preallocator *p = (Mat_Preallocator *) A->data;

 98:   p->nooffproc = PETSC_TRUE;
 99:   while (1) {
100:     MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg);
101:     if (flg) p->nooffproc = PETSC_FALSE;
102:     if (!flg) break;

104:     for (i = 0; i < n;) {
105:       /* Now identify the consecutive vals belonging to the same row */
106:       for (j = i, rstart = row[j]; j < n; j++) {
107:         if (row[j] != rstart) break;
108:       }
109:       if (j < n) ncols = j-i;
110:       else       ncols = n-i;
111:       /* Now assemble all these values with a single function call */
112:       MatSetValues_Preallocator(A, 1, row+i, ncols, col+i, val+i, INSERT_VALUES);
113:       i = j;
114:     }
115:   }
116:   MatStashScatterEnd_Private(&A->stash);
117:   MPIU_Allreduce(MPI_IN_PLACE,&p->nooffproc,1,MPIU_BOOL,MPI_LAND,PetscObjectComm((PetscObject)A));
118:   return 0;
119: }

121: PetscErrorCode MatView_Preallocator(Mat A, PetscViewer viewer)
122: {
123:   return 0;
124: }

126: PetscErrorCode MatSetOption_Preallocator(Mat A, MatOption op, PetscBool flg)
127: {
128:   return 0;
129: }

131: PetscErrorCode MatPreallocatorPreallocate_Preallocator(Mat mat, PetscBool fill, Mat A)
132: {
133:   Mat_Preallocator *p = (Mat_Preallocator *) mat->data;
134:   PetscInt          bs;

137:   p->used = PETSC_TRUE;
138:   if (!fill) PetscHSetIJDestroy(&p->ht);
139:   MatGetBlockSize(mat, &bs);
140:   MatXAIJSetPreallocation(A, bs, p->dnz, p->onz, p->dnzu, p->onzu);
141:   MatSetUp(A);
142:   MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);
143:   MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, p->nooffproc);
144:   if (fill) {
145:     PetscHashIter  hi;
146:     PetscHashIJKey key;
147:     PetscScalar    *zeros;
148:     PetscInt       n,maxrow=1,*cols,rStart,rEnd,*rowstarts;

150:     MatGetOwnershipRange(A, &rStart, &rEnd);
151:     // Ownership range is in terms of scalar entries, but we deal with blocks
152:     rStart /= bs;
153:     rEnd /= bs;
154:     PetscHSetIJGetSize(p->ht,&n);
155:     PetscMalloc2(n,&cols,rEnd-rStart+1,&rowstarts);
156:     rowstarts[0] = 0;
157:     for (PetscInt i=0; i<rEnd-rStart; i++) {
158:       rowstarts[i+1] = rowstarts[i] + p->dnz[i] + p->onz[i];
159:       maxrow = PetscMax(maxrow, p->dnz[i] + p->onz[i]);
160:     }

163:     PetscHashIterBegin(p->ht,hi);
164:     for (PetscInt i=0; !PetscHashIterAtEnd(p->ht,hi); i++) {
165:       PetscHashIterGetKey(p->ht,hi,key);
166:       PetscInt lrow = key.i - rStart;
167:       cols[rowstarts[lrow]] = key.j;
168:       rowstarts[lrow]++;
169:       PetscHashIterNext(p->ht,hi);
170:     }
171:     PetscHSetIJDestroy(&p->ht);

173:     PetscCalloc1(maxrow*bs*bs,&zeros);
174:     for (PetscInt i=0; i<rEnd-rStart; i++) {
175:       PetscInt grow = rStart + i;
176:       PetscInt end = rowstarts[i], start = end - p->dnz[i] - p->onz[i];
177:       PetscSortInt(end-start,&cols[start]);
178:       MatSetValuesBlocked(A, 1, &grow, end-start, &cols[start], zeros, INSERT_VALUES);
179:     }
180:     PetscFree(zeros);
181:     PetscFree2(cols,rowstarts);

183:     MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
184:     MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
185:   }
186:   return 0;
187: }

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

192:   Input Parameters:
193: + mat  - the preallocator
194: . fill - fill the matrix with zeros
195: - A    - the matrix to be preallocated

197:   Notes:
198:   This Mat implementation provides a helper utility to define the correct
199:   preallocation data for a given nonzero structure. Use this object like a
200:   regular matrix, e.g. loop over the nonzero structure of the matrix and
201:   call MatSetValues() or MatSetValuesBlocked() to indicate the nonzero locations.
202:   The matrix entries provided to MatSetValues() will be ignored, it only uses
203:   the row / col indices provided to determine the information required to be
204:   passed to MatXAIJSetPreallocation(). Once you have looped over the nonzero
205:   structure, you must call MatAssemblyBegin(), MatAssemblyEnd() on mat.

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

212:   This function may only be called once for a given MatPreallocator object. If
213:   multiple Mats need to be preallocated, consider using MatDuplicate() after
214:   this function.

216:   Level: advanced

218: .seealso: MATPREALLOCATOR
219: @*/
220: PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A)
221: {
225:   PetscUseMethod(mat,"MatPreallocatorPreallocate_C",(Mat,PetscBool,Mat),(mat,fill,A));
226:   return 0;
227: }

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

232:    Operations Provided:
233: .vb
234:   MatSetValues()
235: .ve

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

240:   Level: advanced

242: .seealso: Mat, MatPreallocatorPreallocate()

244: M*/

246: PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A)
247: {
248:   Mat_Preallocator *p;

250:   PetscNewLog(A, &p);
251:   A->data = (void *) p;

253:   p->ht   = NULL;
254:   p->dnz  = NULL;
255:   p->onz  = NULL;
256:   p->dnzu = NULL;
257:   p->onzu = NULL;
258:   p->used = PETSC_FALSE;

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

263:   A->ops->destroy       = MatDestroy_Preallocator;
264:   A->ops->setup         = MatSetUp_Preallocator;
265:   A->ops->setvalues     = MatSetValues_Preallocator;
266:   A->ops->assemblybegin = MatAssemblyBegin_Preallocator;
267:   A->ops->assemblyend   = MatAssemblyEnd_Preallocator;
268:   A->ops->view          = MatView_Preallocator;
269:   A->ops->setoption     = MatSetOption_Preallocator;
270:   A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */

272:   /* special MATPREALLOCATOR functions */
273:   PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator);
274:   PetscObjectChangeTypeName((PetscObject) A, MATPREALLOCATOR);
275:   return 0;
276: }