1 #include <petsc/private/matimpl.h> /*I "petscmat.h" I*/
2 #include <petsc/private/hashsetij.h>
3
4 typedef struct {
5 PetscHSetIJ ht;
6 PetscInt *dnz, *onz;
7 PetscInt *dnzu, *onzu;
8 PetscBool nooffproc;
9 PetscBool used;
10 } Mat_Preallocator;
11
MatDestroy_Preallocator(Mat A)12 static PetscErrorCode MatDestroy_Preallocator(Mat A)
13 {
14 Mat_Preallocator *p = (Mat_Preallocator *)A->data;
15
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 }
25
MatSetUp_Preallocator(Mat A)26 static PetscErrorCode MatSetUp_Preallocator(Mat A)
27 {
28 Mat_Preallocator *p = (Mat_Preallocator *)A->data;
29 PetscInt m, bs, mbs;
30
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 }
44
MatSetValues_Preallocator(Mat A,PetscInt m,const PetscInt * rows,PetscInt n,const PetscInt * cols,const PetscScalar * values,InsertMode addv)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;
49
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;
57
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 }
82
MatAssemblyBegin_Preallocator(Mat A,MatAssemblyType type)83 static PetscErrorCode MatAssemblyBegin_Preallocator(Mat A, MatAssemblyType type)
84 {
85 PetscInt nstash, reallocs;
86
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 }
93
MatAssemblyEnd_Preallocator(Mat A,MatAssemblyType type)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;
101
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;
108
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, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
123 PetscFunctionReturn(PETSC_SUCCESS);
124 }
125
MatView_Preallocator(Mat A,PetscViewer viewer)126 static PetscErrorCode MatView_Preallocator(Mat A, PetscViewer viewer)
127 {
128 PetscFunctionBegin;
129 PetscFunctionReturn(PETSC_SUCCESS);
130 }
131
MatSetOption_Preallocator(Mat A,MatOption op,PetscBool flg)132 static PetscErrorCode MatSetOption_Preallocator(Mat A, MatOption op, PetscBool flg)
133 {
134 PetscFunctionBegin;
135 PetscFunctionReturn(PETSC_SUCCESS);
136 }
137
MatPreallocatorPreallocate_Preallocator(Mat mat,PetscBool fill,Mat A)138 static PetscErrorCode MatPreallocatorPreallocate_Preallocator(Mat mat, PetscBool fill, Mat A)
139 {
140 Mat_Preallocator *p = (Mat_Preallocator *)mat->data;
141 PetscInt bs;
142
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;
157
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]);
170
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));
180
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));
190
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 }
197
198 /*@
199 MatPreallocatorPreallocate - Preallocates the A matrix, using information from a `MATPREALLOCATOR` mat, optionally filling A with zeros
200
201 Input Parameters:
202 + mat - the `MATPREALLOCATOR` preallocator matrix
203 . fill - fill the matrix with zeros
204 - A - the matrix to be preallocated
205
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.
215
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`);
220
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.
224
225 Level: advanced
226
227 .seealso: `MATPREALLOCATOR`, `MatXAIJSetPreallocation()`
228 @*/
MatPreallocatorPreallocate(Mat mat,PetscBool fill,Mat A)229 PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A)
230 {
231 PetscFunctionBegin;
232 PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
233 PetscValidLogicalCollectiveBool(mat, fill, 2);
234 PetscValidHeaderSpecific(A, MAT_CLASSID, 3);
235 PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat, PetscBool, Mat), (mat, fill, A));
236 PetscFunctionReturn(PETSC_SUCCESS);
237 }
238
239 /*MC
240 MATPREALLOCATOR - MATPREALLOCATOR = "preallocator" - A matrix type to be used for computing a matrix preallocation.
241
242 Operations Provided:
243 .vb
244 MatSetValues()
245 .ve
246
247 Options Database Keys:
248 . -mat_type preallocator - sets the matrix type to `MATPREALLOCATOR` during a call to `MatSetFromOptions()`
249
250 Level: advanced
251
252 .seealso: `MATPREALLOCATOR`, `Mat`, `MatPreallocatorPreallocate()`
253 M*/
254
MatCreate_Preallocator(Mat A)255 PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A)
256 {
257 Mat_Preallocator *p;
258
259 PetscFunctionBegin;
260 PetscCall(PetscNew(&p));
261 A->data = (void *)p;
262
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;
269
270 /* matrix ops */
271 PetscCall(PetscMemzero(A->ops, sizeof(struct _MatOps)));
272
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 */
281
282 /* special MATPREALLOCATOR functions */
283 PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator));
284 PetscCall(PetscObjectChangeTypeName((PetscObject)A, MATPREALLOCATOR));
285 PetscFunctionReturn(PETSC_SUCCESS);
286 }
287