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: static PetscErrorCode MatDestroy_Preallocator(Mat A)
 13: {
 14:   Mat_Preallocator *p = (Mat_Preallocator *)A->data;

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

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

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

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

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

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

 83: static PetscErrorCode MatAssemblyBegin_Preallocator(Mat A, MatAssemblyType type)
 84: {
 85:   PetscInt nstash, reallocs;

 87:   PetscFunctionBegin;
 88:   PetscCall(MatStashScatterBegin_Private(A, &A->stash, A->rmap->range));
 89:   PetscCall(MatStashGetInfo_Private(&A->stash, &nstash, &reallocs));
 90:   PetscCall(PetscInfo(A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
 91:   PetscFunctionReturn(PETSC_SUCCESS);
 92: }

 94: static PetscErrorCode MatAssemblyEnd_Preallocator(Mat A, MatAssemblyType type)
 95: {
 96:   PetscScalar      *val;
 97:   PetscInt         *row, *col;
 98:   PetscInt          i, j, rstart, ncols, flg;
 99:   PetscMPIInt       n;
100:   Mat_Preallocator *p = (Mat_Preallocator *)A->data;

102:   PetscFunctionBegin;
103:   p->nooffproc = PETSC_TRUE;
104:   while (1) {
105:     PetscCall(MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg));
106:     if (flg) p->nooffproc = PETSC_FALSE;
107:     if (!flg) break;

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

126: static PetscErrorCode MatView_Preallocator(Mat A, PetscViewer viewer)
127: {
128:   PetscFunctionBegin;
129:   PetscFunctionReturn(PETSC_SUCCESS);
130: }

132: static PetscErrorCode MatSetOption_Preallocator(Mat A, MatOption op, PetscBool flg)
133: {
134:   PetscFunctionBegin;
135:   PetscFunctionReturn(PETSC_SUCCESS);
136: }

138: static PetscErrorCode MatPreallocatorPreallocate_Preallocator(Mat mat, PetscBool fill, Mat A)
139: {
140:   Mat_Preallocator *p = (Mat_Preallocator *)mat->data;
141:   PetscInt          bs;

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

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

171:     PetscHashIterBegin(p->ht, hi);
172:     while (!PetscHashIterAtEnd(p->ht, hi)) {
173:       PetscHashIterGetKey(p->ht, hi, key);
174:       PetscInt lrow         = key.i - rStart;
175:       cols[rowstarts[lrow]] = key.j;
176:       rowstarts[lrow]++;
177:       PetscHashIterNext(p->ht, hi);
178:     }
179:     PetscCall(PetscHSetIJDestroy(&p->ht));

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

191:     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
192:     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
193:     PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, PETSC_FALSE));
194:   }
195:   PetscFunctionReturn(PETSC_SUCCESS);
196: }

198: /*@
199:   MatPreallocatorPreallocate - Preallocates the A matrix, using information from a `MATPREALLOCATOR` mat, optionally filling A with zeros

201:   Input Parameters:
202: + mat  - the `MATPREALLOCATOR` preallocator matrix
203: . fill - fill the matrix with zeros
204: - A    - the matrix to be preallocated

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

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

221:   This function may only be called once for a given `MATPREALLOCATOR` object. If
222:   multiple `Mat`s need to be preallocated, consider using `MatDuplicate()` after
223:   this function.

225:   Level: advanced

227: .seealso: `MATPREALLOCATOR`, `MatXAIJSetPreallocation()`
228: @*/
229: PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A)
230: {
231:   PetscFunctionBegin;
235:   PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat, PetscBool, Mat), (mat, fill, A));
236:   PetscFunctionReturn(PETSC_SUCCESS);
237: }

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

242:    Operations Provided:
243: .vb
244:   MatSetValues()
245: .ve

247:    Options Database Keys:
248: . -mat_type preallocator - sets the matrix type to `MATPREALLOCATOR` during a call to `MatSetFromOptions()`

250:   Level: advanced

252: .seealso: `MATPREALLOCATOR`, `Mat`, `MatPreallocatorPreallocate()`
253: M*/

255: PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A)
256: {
257:   Mat_Preallocator *p;

259:   PetscFunctionBegin;
260:   PetscCall(PetscNew(&p));
261:   A->data = (void *)p;

263:   p->ht   = NULL;
264:   p->dnz  = NULL;
265:   p->onz  = NULL;
266:   p->dnzu = NULL;
267:   p->onzu = NULL;
268:   p->used = PETSC_FALSE;

270:   /* matrix ops */
271:   PetscCall(PetscMemzero(A->ops, sizeof(struct _MatOps)));

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

282:   /* special MATPREALLOCATOR functions */
283:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator));
284:   PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATPREALLOCATOR));
285:   PetscFunctionReturn(PETSC_SUCCESS);
286: }