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: } 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:   if (!fill) {PetscHSetIJDestroy(&p->ht);}
150:   MatGetBlockSize(mat, &bs);
151:   MatXAIJSetPreallocation(A, bs, p->dnz, p->onz, p->dnzu, p->onzu);
152:   MatSetUp(A);
153:   MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);
154:   MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, p->nooffproc);
155:   if (fill) {
156:     PetscHashIter  hi;
157:     PetscHashIJKey key;
158:     PetscScalar    *zeros;
159:     PetscInt       n,maxrow=1,*cols,rStart,rEnd,*rowstarts;

161:     MatGetOwnershipRange(A, &rStart, &rEnd);
162:     // Ownership range is in terms of scalar entries, but we deal with blocks
163:     rStart /= bs;
164:     rEnd /= bs;
165:     PetscHSetIJGetSize(p->ht,&n);
166:     PetscMalloc2(n,&cols,rEnd-rStart+1,&rowstarts);
167:     rowstarts[0] = 0;
168:     for (PetscInt i=0; i<rEnd-rStart; i++) {
169:       rowstarts[i+1] = rowstarts[i] + p->dnz[i] + p->onz[i];
170:       maxrow = PetscMax(maxrow, p->dnz[i] + p->onz[i]);
171:     }
172:     if (rowstarts[rEnd-rStart] != n) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Hash claims %D entries, but dnz+onz counts %D",n,rowstarts[rEnd-rStart]);

174:     PetscHashIterBegin(p->ht,hi);
175:     for (PetscInt i=0; !PetscHashIterAtEnd(p->ht,hi); i++) {
176:       PetscHashIterGetKey(p->ht,hi,key);
177:       PetscInt lrow = key.i - rStart;
178:       cols[rowstarts[lrow]] = key.j;
179:       rowstarts[lrow]++;
180:       PetscHashIterNext(p->ht,hi);
181:     }
182:     PetscHSetIJDestroy(&p->ht);

184:     PetscCalloc1(maxrow*bs*bs,&zeros);
185:     for (PetscInt i=0; i<rEnd-rStart; i++) {
186:       PetscInt grow = rStart + i;
187:       PetscInt end = rowstarts[i], start = end - p->dnz[i] - p->onz[i];
188:       PetscSortInt(end-start,&cols[start]);
189:       MatSetValuesBlocked(A, 1, &grow, end-start, &cols[start], zeros, INSERT_VALUES);
190:     }
191:     PetscFree(zeros);
192:     PetscFree2(cols,rowstarts);

194:     MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
195:     MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
196:   }
197:   return(0);
198: }

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

203:   Input Parameters:
204: + mat  - the preallocator
205: . fill - fill the matrix with zeros
206: - A    - the matrix to be preallocated

208:   Notes:
209:   This Mat implementation provides a helper utility to define the correct
210:   preallocation data for a given nonzero structure. Use this object like a
211:   regular matrix, e.g. loop over the nonzero structure of the matrix and
212:   call MatSetValues() or MatSetValuesBlocked() to indicate the nonzero locations.
213:   The matrix entries provided to MatSetValues() will be ignored, it only uses
214:   the row / col indices provided to determine the information required to be
215:   passed to MatXAIJSetPreallocation(). Once you have looped over the nonzero
216:   structure, you must call MatAssemblyBegin(), MatAssemblyEnd() on mat.

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

223:   Level: advanced

225: .seealso: MATPREALLOCATOR
226: @*/
227: PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A)
228: {

234:   PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat,PetscBool,Mat),(mat,fill,A));
235:   return(0);
236: }

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

241:    Operations Provided:
242: .  MatSetValues()

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

247:   Level: advanced

249: .seealso: Mat

251: M*/

253: PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A)
254: {
255:   Mat_Preallocator *p;
256:   PetscErrorCode    ierr;

259:   PetscNewLog(A, &p);
260:   A->data = (void *) p;

262:   p->ht   = NULL;
263:   p->dnz  = NULL;
264:   p->onz  = NULL;
265:   p->dnzu = NULL;
266:   p->onzu = NULL;

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

271:   A->ops->destroy       = MatDestroy_Preallocator;
272:   A->ops->setup         = MatSetUp_Preallocator;
273:   A->ops->setvalues     = MatSetValues_Preallocator;
274:   A->ops->assemblybegin = MatAssemblyBegin_Preallocator;
275:   A->ops->assemblyend   = MatAssemblyEnd_Preallocator;
276:   A->ops->view          = MatView_Preallocator;
277:   A->ops->setoption     = MatSetOption_Preallocator;
278:   A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */

280:   /* special MATPREALLOCATOR functions */
281:   PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator);
282:   PetscObjectChangeTypeName((PetscObject) A, MATPREALLOCATOR);
283:   return(0);
284: }