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