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;
 15:   PetscErrorCode    ierr;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

150:   if (p->used) SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"MatPreallocatorPreallocate() can only be used once for a give MatPreallocator object. Consider using MatDuplicate() after preallocation.");
151:   p->used = PETSC_TRUE;
152:   if (!fill) {PetscHSetIJDestroy(&p->ht);}
153:   MatGetBlockSize(mat, &bs);
154:   MatXAIJSetPreallocation(A, bs, p->dnz, p->onz, p->dnzu, p->onzu);
155:   MatSetUp(A);
156:   MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);
157:   MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, p->nooffproc);
158:   if (fill) {
159:     PetscHashIter  hi;
160:     PetscHashIJKey key;
161:     PetscScalar    *zeros;
162:     PetscInt       n,maxrow=1,*cols,rStart,rEnd,*rowstarts;

164:     MatGetOwnershipRange(A, &rStart, &rEnd);
165:     // Ownership range is in terms of scalar entries, but we deal with blocks
166:     rStart /= bs;
167:     rEnd /= bs;
168:     PetscHSetIJGetSize(p->ht,&n);
169:     PetscMalloc2(n,&cols,rEnd-rStart+1,&rowstarts);
170:     rowstarts[0] = 0;
171:     for (PetscInt i=0; i<rEnd-rStart; i++) {
172:       rowstarts[i+1] = rowstarts[i] + p->dnz[i] + p->onz[i];
173:       maxrow = PetscMax(maxrow, p->dnz[i] + p->onz[i]);
174:     }
175:     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]);

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

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

197:     MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
198:     MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
199:   }
200:   return(0);
201: }

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

206:   Input Parameters:
207: + mat  - the preallocator
208: . fill - fill the matrix with zeros
209: - A    - the matrix to be preallocated

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

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

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

230:   Level: advanced

232: .seealso: MATPREALLOCATOR
233: @*/
234: PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A)
235: {

241:   PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat,PetscBool,Mat),(mat,fill,A));
242:   return(0);
243: }

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

248:    Operations Provided:
249: .  MatSetValues()

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

254:   Level: advanced

256: .seealso: Mat, MatPreallocatorPreallocate()

258: M*/

260: PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A)
261: {
262:   Mat_Preallocator *p;
263:   PetscErrorCode    ierr;

266:   PetscNewLog(A, &p);
267:   A->data = (void *) p;

269:   p->ht   = NULL;
270:   p->dnz  = NULL;
271:   p->onz  = NULL;
272:   p->dnzu = NULL;
273:   p->onzu = NULL;
274:   p->used = PETSC_FALSE;

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

279:   A->ops->destroy       = MatDestroy_Preallocator;
280:   A->ops->setup         = MatSetUp_Preallocator;
281:   A->ops->setvalues     = MatSetValues_Preallocator;
282:   A->ops->assemblybegin = MatAssemblyBegin_Preallocator;
283:   A->ops->assemblyend   = MatAssemblyEnd_Preallocator;
284:   A->ops->view          = MatView_Preallocator;
285:   A->ops->setoption     = MatSetOption_Preallocator;
286:   A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */

288:   /* special MATPREALLOCATOR functions */
289:   PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator);
290:   PetscObjectChangeTypeName((PetscObject) A, MATPREALLOCATOR);
291:   return(0);
292: }