xref: /petsc/src/mat/impls/preallocator/matpreallocator.c (revision 2da392cc7c10228af19ad9843ce5155178acb644)
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 } Mat_Preallocator;
9 
10 PetscErrorCode MatDestroy_Preallocator(Mat A)
11 {
12   Mat_Preallocator *p = (Mat_Preallocator *) A->data;
13   PetscErrorCode    ierr;
14 
15   PetscFunctionBegin;
16   ierr = MatStashDestroy_Private(&A->stash);CHKERRQ(ierr);
17   ierr = PetscHSetIJDestroy(&p->ht);CHKERRQ(ierr);
18   ierr = PetscFree4(p->dnz, p->onz, p->dnzu, p->onzu);CHKERRQ(ierr);
19   ierr = PetscFree(A->data);CHKERRQ(ierr);
20   ierr = PetscObjectChangeTypeName((PetscObject) A, NULL);CHKERRQ(ierr);
21   ierr = PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", NULL);CHKERRQ(ierr);
22   PetscFunctionReturn(0);
23 }
24 
25 PetscErrorCode MatSetUp_Preallocator(Mat A)
26 {
27   Mat_Preallocator *p = (Mat_Preallocator *) A->data;
28   PetscInt          m, bs, mbs;
29   PetscErrorCode    ierr;
30 
31   PetscFunctionBegin;
32   ierr = PetscLayoutSetUp(A->rmap);CHKERRQ(ierr);
33   ierr = PetscLayoutSetUp(A->cmap);CHKERRQ(ierr);
34   ierr = MatGetLocalSize(A, &m, NULL);CHKERRQ(ierr);
35   ierr = PetscHSetIJCreate(&p->ht);CHKERRQ(ierr);
36   ierr = MatGetBlockSize(A, &bs);CHKERRQ(ierr);
37   /* Do not bother bstash since MatPreallocator does not implement MatSetValuesBlocked */
38   ierr = MatStashCreate_Private(PetscObjectComm((PetscObject) A), 1, &A->stash);CHKERRQ(ierr);
39   /* arrays are for blocked rows/cols */
40   mbs  = m/bs;
41   ierr = PetscCalloc4(mbs, &p->dnz, mbs, &p->onz, mbs, &p->dnzu, mbs, &p->onzu);CHKERRQ(ierr);
42   PetscFunctionReturn(0);
43 }
44 
45 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   PetscErrorCode    ierr;
50 
51   PetscFunctionBegin;
52   ierr = MatGetBlockSize(A, &bs);CHKERRQ(ierr);
53   ierr = MatGetOwnershipRange(A, &rStart, &rEnd);CHKERRQ(ierr);
54   ierr = MatGetOwnershipRangeColumn(A, &cStart, &cEnd);CHKERRQ(ierr);
55   for (r = 0; r < m; ++r) {
56     PetscHashIJKey key;
57     PetscBool      missing;
58 
59     key.i = rows[r];
60     if (key.i < 0) continue;
61     if ((key.i < rStart) || (key.i >= rEnd)) {
62       ierr = MatStashValuesRow_Private(&A->stash, key.i, n, cols, values, PETSC_FALSE);CHKERRQ(ierr);
63     } else { /* Hash table is for blocked rows/cols */
64       key.i = rows[r]/bs;
65       for (c = 0; c < n; ++c) {
66         key.j = cols[c]/bs;
67         if (key.j < 0) continue;
68         ierr = PetscHSetIJQueryAdd(p->ht, key, &missing);CHKERRQ(ierr);
69         if (missing) {
70           if ((key.j >= cStart/bs) && (key.j < cEnd/bs)) {
71             ++p->dnz[key.i-rStart/bs];
72             if (key.j >= key.i) ++p->dnzu[key.i-rStart/bs];
73           } else {
74             ++p->onz[key.i-rStart/bs];
75             if (key.j >= key.i) ++p->onzu[key.i-rStart/bs];
76           }
77         }
78       }
79     }
80   }
81   PetscFunctionReturn(0);
82 }
83 
84 PetscErrorCode MatAssemblyBegin_Preallocator(Mat A, MatAssemblyType type)
85 {
86   PetscInt       nstash, reallocs;
87   PetscErrorCode ierr;
88 
89   PetscFunctionBegin;
90   ierr = MatStashScatterBegin_Private(A, &A->stash, A->rmap->range);CHKERRQ(ierr);
91   ierr = MatStashGetInfo_Private(&A->stash, &nstash, &reallocs);CHKERRQ(ierr);
92   ierr = PetscInfo2(A, "Stash has %D entries, uses %D mallocs.\n", nstash, reallocs);CHKERRQ(ierr);
93   PetscFunctionReturn(0);
94 }
95 
96 PetscErrorCode MatAssemblyEnd_Preallocator(Mat A, MatAssemblyType type)
97 {
98   PetscScalar   *val;
99   PetscInt      *row, *col;
100   PetscInt       i, j, rstart, ncols, flg;
101   PetscMPIInt    n;
102   PetscErrorCode ierr;
103 
104   PetscFunctionBegin;
105   while (1) {
106     ierr = MatStashScatterGetMesg_Private(&A->stash, &n, &row, &col, &val, &flg);CHKERRQ(ierr);
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       ierr = MatSetValues_Preallocator(A, 1, row+i, ncols, col+i, val+i, INSERT_VALUES);CHKERRQ(ierr);
118       i = j;
119     }
120   }
121   ierr = MatStashScatterEnd_Private(&A->stash);CHKERRQ(ierr);
122   PetscFunctionReturn(0);
123 }
124 
125 PetscErrorCode MatView_Preallocator(Mat A, PetscViewer viewer)
126 {
127   PetscFunctionBegin;
128   PetscFunctionReturn(0);
129 }
130 
131 PetscErrorCode MatSetOption_Preallocator(Mat A, MatOption op, PetscBool flg)
132 {
133   PetscFunctionBegin;
134   PetscFunctionReturn(0);
135 }
136 
137 PetscErrorCode MatPreallocatorPreallocate_Preallocator(Mat mat, PetscBool fill, Mat A)
138 {
139   Mat_Preallocator *p = (Mat_Preallocator *) mat->data;
140   PetscInt          bs;
141   PetscErrorCode    ierr;
142 
143   PetscFunctionBegin;
144   ierr = MatGetBlockSize(mat, &bs);CHKERRQ(ierr);
145   ierr = MatXAIJSetPreallocation(A, bs, p->dnz, p->onz, p->dnzu, p->onzu);CHKERRQ(ierr);
146   ierr = MatSetUp(A);CHKERRQ(ierr);
147   ierr = MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);CHKERRQ(ierr);
148   if (fill) {
149     PetscHashIter  hi;
150     PetscHashIJKey key;
151     PetscScalar    *zeros;
152 
153     ierr = PetscCalloc1(bs*bs,&zeros);CHKERRQ(ierr);
154 
155     PetscHashIterBegin(p->ht,hi);
156     while (!PetscHashIterAtEnd(p->ht,hi)) {
157       PetscHashIterGetKey(p->ht,hi,key);
158       PetscHashIterNext(p->ht,hi);
159       ierr = MatSetValuesBlocked(A,1,&key.i,1,&key.j,zeros,INSERT_VALUES);CHKERRQ(ierr);
160     }
161     ierr = PetscFree(zeros);CHKERRQ(ierr);
162 
163     ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
164     ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
165   }
166   PetscFunctionReturn(0);
167 }
168 
169 /*@
170   MatPreallocatorPreallocate - Preallocates the A matrix, using information from mat, optionally filling A with zeros
171 
172   Input Parameter:
173 + mat  - the preallocator
174 . fill - fill the matrix with zeros
175 - A    - the matrix to be preallocated
176 
177   Notes:
178   This Mat implementation provides a helper utility to define the correct
179   preallocation data for a given nonzero structure. Use this object like a
180   regular matrix, e.g. loop over the nonzero structure of the matrix and
181   call MatSetValues() or MatSetValuesBlocked() to indicate the nonzero locations.
182   The matrix entires provided to MatSetValues() will be ignored, it only uses
183   the row / col indices provided to determine the information required to be
184   passed to MatXAIJSetPreallocation(). Once you have looped over the nonzero
185   structure, you must call MatAssemblyBegin(), MatAssemblyEnd() on mat.
186 
187   After you have assembled the preallocator matrix (mat), call MatPreallocatorPreallocate()
188   to define the preallocation information on the matrix (A). Setting the parameter
189   fill = PETSC_TRUE will insert zeros into the matrix A. Internally MatPreallocatorPreallocate()
190   will call MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);
191 
192   Level: advanced
193 
194 .seealso: MATPREALLOCATOR
195 @*/
196 PetscErrorCode MatPreallocatorPreallocate(Mat mat, PetscBool fill, Mat A)
197 {
198   PetscErrorCode ierr;
199 
200   PetscFunctionBegin;
201   PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
202   PetscValidHeaderSpecific(A,   MAT_CLASSID, 3);
203   ierr = PetscUseMethod(mat, "MatPreallocatorPreallocate_C", (Mat,PetscBool,Mat),(mat,fill,A));CHKERRQ(ierr);
204   PetscFunctionReturn(0);
205 }
206 
207 /*MC
208    MATPREALLOCATOR - MATPREALLOCATOR = "preallocator" - A matrix type to be used for computing a matrix preallocation.
209 
210    Operations Provided:
211 .  MatSetValues()
212 
213    Options Database Keys:
214 . -mat_type preallocator - sets the matrix type to "preallocator" during a call to MatSetFromOptions()
215 
216   Level: advanced
217 
218 .seealso: Mat
219 
220 M*/
221 
222 PETSC_EXTERN PetscErrorCode MatCreate_Preallocator(Mat A)
223 {
224   Mat_Preallocator *p;
225   PetscErrorCode    ierr;
226 
227   PetscFunctionBegin;
228   ierr = PetscNewLog(A, &p);CHKERRQ(ierr);
229   A->data = (void *) p;
230 
231   p->ht   = NULL;
232   p->dnz  = NULL;
233   p->onz  = NULL;
234   p->dnzu = NULL;
235   p->onzu = NULL;
236 
237   /* matrix ops */
238   ierr = PetscMemzero(A->ops, sizeof(struct _MatOps));CHKERRQ(ierr);
239 
240   A->ops->destroy       = MatDestroy_Preallocator;
241   A->ops->setup         = MatSetUp_Preallocator;
242   A->ops->setvalues     = MatSetValues_Preallocator;
243   A->ops->assemblybegin = MatAssemblyBegin_Preallocator;
244   A->ops->assemblyend   = MatAssemblyEnd_Preallocator;
245   A->ops->view          = MatView_Preallocator;
246   A->ops->setoption     = MatSetOption_Preallocator;
247   A->ops->setblocksizes = MatSetBlockSizes_Default; /* once set, user is not allowed to change the block sizes */
248 
249   /* special MATPREALLOCATOR functions */
250   ierr = PetscObjectComposeFunction((PetscObject) A, "MatPreallocatorPreallocate_C", MatPreallocatorPreallocate_Preallocator);CHKERRQ(ierr);
251   ierr = PetscObjectChangeTypeName((PetscObject) A, MATPREALLOCATOR);CHKERRQ(ierr);
252   PetscFunctionReturn(0);
253 }
254