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