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: PetscCallMPI(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, PetscSafePointerPlusOffset(cols, start)));
186: PetscCall(MatSetValuesBlocked(A, 1, &grow, end - start, PetscSafePointerPlusOffset(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: }