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: if (fill) {
144: MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, p->nooffproc);
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: while (!PetscHashIterAtEnd(p->ht, hi)) {
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: MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, PETSC_FALSE);
186: }
187: return 0;
188: }
190: /*@
191: MatPreallocatorPreallocate - Preallocates the A matrix, using information from a `MATPREALLOCATOR` mat, optionally filling A with zeros
193: Input Parameters:
194: + mat - the `MATPREALLOCATOR` preallocator matrix
195: . fill - fill the matrix with zeros
196: - A - the matrix to be preallocated
198: Notes:
199: This `MatType` implementation provides a helper utility to define the correct
200: preallocation data for a given nonzero structure. Use this object like a
201: regular matrix, e.g. loop over the nonzero structure of the matrix and
202: call `MatSetValues()` or `MatSetValuesBlocked()` to indicate the nonzero locations.
203: The matrix entries provided to `MatSetValues()` will be ignored, it only uses
204: the row / col indices provided to determine the information required to be
205: passed to `MatXAIJSetPreallocation()`. Once you have looped over the nonzero
206: structure, you must call `MatAssemblyBegin()`, `MatAssemblyEnd()` on mat.
208: After you have assembled the preallocator matrix (mat), call `MatPreallocatorPreallocate()`
209: to define the preallocation information on the matrix (A). Setting the parameter
210: fill = PETSC_TRUE will insert zeros into the matrix A. Internally `MatPreallocatorPreallocate()`
211: will call `MatSetOption`(A, `MAT_NEW_NONZERO_ALLOCATION_ERR`, `PETSC_TRUE`);
213: This function may only be called once for a given `MATPREALLOCATOR` object. If
214: multiple `Mat`s need to be preallocated, consider using `MatDuplicate()` after
215: this function.
217: Level: advanced
219: .seealso: `MATPREALLOCATOR`, `MatXAIJSetPreallocation()`
220: @*/
221: PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A)
222: {
226: PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat, PetscBool, Mat), (mat, fill, A));
227: return 0;
228: }
230: /*MC
231: MATPREALLOCATOR - MATPREALLOCATOR = "preallocator" - A matrix type to be used for computing a matrix preallocation.
233: Operations Provided:
234: .vb
235: MatSetValues()
236: .ve
238: Options Database Keys:
239: . -mat_type preallocator - sets the matrix type to `MATPREALLOCATOR` during a call to `MatSetFromOptions()`
241: Level: advanced
243: .seealso: `MATPREALLOCATOR`, `Mat`, `MatPreallocatorPreallocate()`
244: M*/
246: PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A)
247: {
248: Mat_Preallocator *p;
250: PetscNew(&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: }