xref: /petsc/src/mat/impls/sbaij/seq/aijsbaij.c (revision bcd4bb4a4158aa96f212e9537e87b40407faf83e)
1 #include <../src/mat/impls/aij/seq/aij.h>
2 #include <../src/mat/impls/baij/seq/baij.h>
3 #include <../src/mat/impls/sbaij/seq/sbaij.h>
4 
5 PETSC_INTERN PetscErrorCode MatConvert_SeqSBAIJ_SeqAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
6 {
7   Mat           B;
8   Mat_SeqSBAIJ *a = (Mat_SeqSBAIJ *)A->data;
9   Mat_SeqAIJ   *b;
10   PetscInt     *ai = a->i, *aj = a->j, m = A->rmap->N, n = A->cmap->n, i, j, k, *bi, *bj, *rowlengths, nz, *rowstart, itmp;
11   PetscInt      bs = A->rmap->bs, bs2 = bs * bs, mbs = A->rmap->N / bs, diagcnt = 0;
12   MatScalar    *av, *bv;
13 #if defined(PETSC_USE_COMPLEX)
14   const int aconj = A->hermitian == PETSC_BOOL3_TRUE ? 1 : 0;
15 #else
16   const int aconj = 0;
17 #endif
18 
19   PetscFunctionBegin;
20   /* compute rowlengths of newmat */
21   PetscCall(PetscMalloc2(m, &rowlengths, m + 1, &rowstart));
22 
23   for (i = 0; i < mbs; i++) rowlengths[i * bs] = 0;
24   k = 0;
25   for (i = 0; i < mbs; i++) {
26     nz = ai[i + 1] - ai[i];
27     if (nz) {
28       rowlengths[k] += nz; /* no. of upper triangular blocks */
29       if (*aj == i) {
30         aj++;
31         diagcnt++;
32         nz--;
33       } /* skip diagonal */
34       for (j = 0; j < nz; j++) { /* no. of lower triangular blocks */
35         rowlengths[(*aj) * bs]++;
36         aj++;
37       }
38     }
39     rowlengths[k] *= bs;
40     for (j = 1; j < bs; j++) rowlengths[k + j] = rowlengths[k];
41     k += bs;
42   }
43 
44   if (reuse != MAT_REUSE_MATRIX) {
45     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
46     PetscCall(MatSetSizes(B, m, n, m, n));
47     PetscCall(MatSetType(B, MATSEQAIJ));
48     PetscCall(MatSeqAIJSetPreallocation(B, 0, rowlengths));
49     PetscCall(MatSetBlockSize(B, A->rmap->bs));
50   } else B = *newmat;
51 
52   b  = (Mat_SeqAIJ *)B->data;
53   bi = b->i;
54   bj = b->j;
55   bv = b->a;
56 
57   /* set b->i */
58   bi[0]       = 0;
59   rowstart[0] = 0;
60   for (i = 0; i < mbs; i++) {
61     for (j = 0; j < bs; j++) {
62       b->ilen[i * bs + j]      = rowlengths[i * bs];
63       rowstart[i * bs + j + 1] = rowstart[i * bs + j] + rowlengths[i * bs];
64     }
65     bi[i + 1] = bi[i] + rowlengths[i * bs] / bs;
66   }
67   PetscCheck(bi[mbs] == 2 * a->nz - diagcnt, PETSC_COMM_SELF, PETSC_ERR_PLIB, "bi[mbs]: %" PetscInt_FMT " != 2*a->nz-diagcnt: %" PetscInt_FMT, bi[mbs], 2 * a->nz - diagcnt);
68 
69   /* set b->j and b->a */
70   aj = a->j;
71   av = a->a;
72   for (i = 0; i < mbs; i++) {
73     nz = ai[i + 1] - ai[i];
74     /* diagonal block */
75     if (nz && *aj == i) {
76       nz--;
77       for (j = 0; j < bs; j++) { /* row i*bs+j */
78         itmp = i * bs + j;
79         for (k = 0; k < bs; k++) { /* col i*bs+k */
80           *(bj + rowstart[itmp]) = (*aj) * bs + k;
81           *(bv + rowstart[itmp]) = *(av + k * bs + j);
82           rowstart[itmp]++;
83         }
84       }
85       aj++;
86       av += bs2;
87     }
88 
89     while (nz--) {
90       /* lower triangular blocks */
91       for (j = 0; j < bs; j++) { /* row (*aj)*bs+j */
92         itmp = (*aj) * bs + j;
93         for (k = 0; k < bs; k++) { /* col i*bs+k */
94           *(bj + rowstart[itmp]) = i * bs + k;
95           *(bv + rowstart[itmp]) = aconj ? PetscConj(*(av + j * bs + k)) : *(av + j * bs + k);
96           rowstart[itmp]++;
97         }
98       }
99       /* upper triangular blocks */
100       for (j = 0; j < bs; j++) { /* row i*bs+j */
101         itmp = i * bs + j;
102         for (k = 0; k < bs; k++) { /* col (*aj)*bs+k */
103           *(bj + rowstart[itmp]) = (*aj) * bs + k;
104           *(bv + rowstart[itmp]) = *(av + k * bs + j);
105           rowstart[itmp]++;
106         }
107       }
108       aj++;
109       av += bs2;
110     }
111   }
112   PetscCall(PetscFree2(rowlengths, rowstart));
113   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
114   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
115 
116   if (reuse == MAT_INPLACE_MATRIX) {
117     PetscCall(MatHeaderReplace(A, &B));
118   } else {
119     *newmat = B;
120   }
121   PetscFunctionReturn(PETSC_SUCCESS);
122 }
123 
124 #include <../src/mat/impls/aij/seq/aij.h>
125 
126 PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqSBAIJ_Preallocate(Mat A, PetscInt **nnz)
127 {
128   Mat_SeqAIJ     *Aa = (Mat_SeqAIJ *)A->data;
129   PetscInt        m, n, bs = A->rmap->bs;
130   const PetscInt *ai = Aa->i, *aj = Aa->j;
131 
132   PetscFunctionBegin;
133   PetscCall(MatGetSize(A, &m, &n));
134 
135   if (bs == 1) {
136     const PetscInt *adiag;
137 
138     PetscCall(MatGetDiagonalMarkers_SeqAIJ(A, &adiag, NULL));
139     PetscCall(PetscMalloc1(m, nnz));
140     for (PetscInt i = 0; i < m; i++) {
141       if (adiag[i] == ai[i + 1]) {
142         (*nnz)[i] = 0;
143         for (PetscInt j = ai[i]; j < ai[i + 1]; j++) (*nnz)[i] += (aj[j] > i);
144       } else (*nnz)[i] = ai[i + 1] - adiag[i];
145     }
146   } else {
147     PetscHSetIJ    ht;
148     PetscHashIJKey key;
149     PetscBool      missing;
150 
151     PetscCall(PetscHSetIJCreate(&ht));
152     PetscCall(PetscCalloc1(m / bs, nnz));
153     for (PetscInt i = 0; i < m; i++) {
154       key.i = i / bs;
155       for (PetscInt k = ai[i]; k < ai[i + 1]; k++) {
156         key.j = aj[k] / bs;
157         if (key.j < key.i) continue;
158         PetscCall(PetscHSetIJQueryAdd(ht, key, &missing));
159         if (missing) (*nnz)[key.i]++;
160       }
161     }
162     PetscCall(PetscHSetIJDestroy(&ht));
163   }
164   PetscFunctionReturn(PETSC_SUCCESS);
165 }
166 
167 PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqSBAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
168 {
169   Mat             B;
170   Mat_SeqAIJ     *a = (Mat_SeqAIJ *)A->data;
171   Mat_SeqSBAIJ   *b;
172   PetscInt       *ai = a->i, *aj, m = A->rmap->N, n = A->cmap->N, i, j, *bi, *bj, *rowlengths, bs = A->rmap->bs;
173   MatScalar      *av, *bv;
174   PetscBool       miss = PETSC_FALSE;
175   const PetscInt *adiag;
176 
177   PetscFunctionBegin;
178 #if !defined(PETSC_USE_COMPLEX)
179   PetscCheck(A->symmetric == PETSC_BOOL3_TRUE, PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Matrix must be symmetric. Call MatSetOption(mat,MAT_SYMMETRIC,PETSC_TRUE)");
180 #else
181   PetscCheck(A->symmetric == PETSC_BOOL3_TRUE || A->hermitian == PETSC_BOOL3_TRUE, PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Matrix must be either symmetric or Hermitian. Call MatSetOption(mat,MAT_SYMMETRIC,PETSC_TRUE) and/or MatSetOption(mat,MAT_HERMITIAN,PETSC_TRUE)");
182 #endif
183   PetscCheck(n == m, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Matrix must be square");
184   PetscCall(MatGetDiagonalMarkers_SeqAIJ(A, &adiag, NULL));
185   if (bs == 1) {
186     PetscCall(PetscMalloc1(m, &rowlengths));
187     for (i = 0; i < m; i++) {
188       if (adiag[i] == ai[i + 1]) {               /* missing diagonal */
189         rowlengths[i] = (ai[i + 1] - ai[i]) + 1; /* allocate some extra space */
190         miss          = PETSC_TRUE;
191       } else {
192         rowlengths[i] = (ai[i + 1] - adiag[i]);
193       }
194     }
195   } else PetscCall(MatConvert_SeqAIJ_SeqSBAIJ_Preallocate(A, &rowlengths));
196   if (reuse != MAT_REUSE_MATRIX) {
197     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
198     PetscCall(MatSetSizes(B, m, n, m, n));
199     PetscCall(MatSetType(B, MATSEQSBAIJ));
200     PetscCall(MatSeqSBAIJSetPreallocation(B, bs, 0, rowlengths));
201   } else B = *newmat;
202 
203   if (bs == 1 && !miss) {
204     b  = (Mat_SeqSBAIJ *)B->data;
205     bi = b->i;
206     bj = b->j;
207     bv = b->a;
208 
209     bi[0] = 0;
210     for (i = 0; i < m; i++) {
211       aj = a->j + adiag[i];
212       av = a->a + adiag[i];
213       for (j = 0; j < rowlengths[i]; j++) {
214         *bj = *aj;
215         bj++;
216         aj++;
217         *bv = *av;
218         bv++;
219         av++;
220       }
221       bi[i + 1]  = bi[i] + rowlengths[i];
222       b->ilen[i] = rowlengths[i];
223     }
224     PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
225     PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
226   } else {
227     PetscCall(MatSetOption(B, MAT_IGNORE_LOWER_TRIANGULAR, PETSC_TRUE));
228     /* reuse may not be equal to MAT_REUSE_MATRIX, but the basic converter will reallocate or replace newmat if this value is not used */
229     /* if reuse is equal to MAT_INITIAL_MATRIX, it has been appropriately preallocated before                                          */
230     /*                      MAT_INPLACE_MATRIX, it will be replaced with MatHeaderReplace below                                        */
231     PetscCall(MatConvert_Basic(A, newtype, MAT_REUSE_MATRIX, &B));
232   }
233   PetscCall(PetscFree(rowlengths));
234   if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &B));
235   else *newmat = B;
236   PetscFunctionReturn(PETSC_SUCCESS);
237 }
238 
239 PETSC_INTERN PetscErrorCode MatConvert_SeqSBAIJ_SeqBAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
240 {
241   Mat           B;
242   Mat_SeqSBAIJ *a = (Mat_SeqSBAIJ *)A->data;
243   Mat_SeqBAIJ  *b;
244   PetscInt     *ai = a->i, *aj = a->j, m = A->rmap->N, n = A->cmap->n, *bi, *bj, *browlengths, nz, *browstart, itmp;
245   PetscInt      bs = A->rmap->bs, bs2 = bs * bs, mbs = m / bs;
246   MatScalar    *av, *bv;
247 #if defined(PETSC_USE_COMPLEX)
248   const int aconj = A->hermitian == PETSC_BOOL3_TRUE ? 1 : 0;
249 #else
250   const int aconj = 0;
251 #endif
252 
253   PetscFunctionBegin;
254   /* compute browlengths of newmat */
255   PetscCall(PetscMalloc2(mbs, &browlengths, mbs, &browstart));
256   for (PetscInt i = 0; i < mbs; i++) browlengths[i] = 0;
257   for (PetscInt i = 0; i < mbs; i++) {
258     nz = ai[i + 1] - ai[i];
259     aj++;                               /* skip diagonal */
260     for (PetscInt k = 1; k < nz; k++) { /* no. of lower triangular blocks */
261       browlengths[*aj]++;
262       aj++;
263     }
264     browlengths[i] += nz; /* no. of upper triangular blocks */
265   }
266 
267   if (reuse != MAT_REUSE_MATRIX) {
268     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
269     PetscCall(MatSetSizes(B, m, n, m, n));
270     PetscCall(MatSetType(B, MATSEQBAIJ));
271     PetscCall(MatSeqBAIJSetPreallocation(B, bs, 0, browlengths));
272   } else B = *newmat;
273 
274   b  = (Mat_SeqBAIJ *)B->data;
275   bi = b->i;
276   bj = b->j;
277   bv = b->a;
278 
279   /* set b->i */
280   bi[0] = 0;
281   for (PetscInt i = 0; i < mbs; i++) {
282     b->ilen[i]   = browlengths[i];
283     bi[i + 1]    = bi[i] + browlengths[i];
284     browstart[i] = bi[i];
285   }
286   PetscCheck(bi[mbs] == 2 * a->nz - mbs, PETSC_COMM_SELF, PETSC_ERR_PLIB, "bi[mbs]: %" PetscInt_FMT " != 2*a->nz - mbs: %" PetscInt_FMT, bi[mbs], 2 * a->nz - mbs);
287 
288   /* set b->j and b->a */
289   aj = a->j;
290   av = a->a;
291   for (PetscInt i = 0; i < mbs; i++) {
292     /* diagonal block */
293     *(bj + browstart[i]) = *aj;
294     aj++;
295 
296     itmp = bs2 * browstart[i];
297     for (PetscInt k = 0; k < bs2; k++) {
298       *(bv + itmp + k) = *av;
299       av++;
300     }
301     browstart[i]++;
302 
303     nz = ai[i + 1] - ai[i] - 1;
304     while (nz--) {
305       /* lower triangular blocks - transpose blocks of A */
306       *(bj + browstart[*aj]) = i; /* block col index */
307 
308       itmp = bs2 * browstart[*aj]; /* row index */
309       for (PetscInt col = 0; col < bs; col++) {
310         PetscInt k = col;
311         for (PetscInt row = 0; row < bs; row++) {
312           bv[itmp + col * bs + row] = aconj ? PetscConj(av[k]) : av[k];
313           k += bs;
314         }
315       }
316       browstart[*aj]++;
317 
318       /* upper triangular blocks */
319       *(bj + browstart[i]) = *aj;
320       aj++;
321 
322       itmp = bs2 * browstart[i];
323       for (PetscInt k = 0; k < bs2; k++) bv[itmp + k] = av[k];
324       av += bs2;
325       browstart[i]++;
326     }
327   }
328   PetscCall(PetscFree2(browlengths, browstart));
329   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
330   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
331 
332   if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &B));
333   else *newmat = B;
334   PetscFunctionReturn(PETSC_SUCCESS);
335 }
336 
337 PETSC_INTERN PetscErrorCode MatConvert_SeqBAIJ_SeqSBAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
338 {
339   Mat             B;
340   Mat_SeqBAIJ    *a = (Mat_SeqBAIJ *)A->data;
341   Mat_SeqSBAIJ   *b;
342   PetscInt       *ai = a->i, *aj, m = A->rmap->N, n = A->cmap->n, k, *bi, *bj, *browlengths;
343   PetscInt        bs = A->rmap->bs, bs2 = bs * bs, mbs = m / bs;
344   MatScalar      *av, *bv;
345   const PetscInt *adiag;
346 
347   PetscFunctionBegin;
348   PetscCheck(A->symmetric == PETSC_BOOL3_TRUE, PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Matrix must be symmetric. Call MatSetOption(mat,MAT_SYMMETRIC,PETSC_TRUE)");
349   PetscCheck(n == m, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Matrix must be square");
350   PetscCall(MatGetDiagonalMarkers_SeqBAIJ(A, &adiag, NULL));
351   PetscCall(PetscMalloc1(mbs, &browlengths));
352   for (PetscInt i = 0; i < mbs; i++) browlengths[i] = ai[i + 1] - adiag[i];
353 
354   if (reuse != MAT_REUSE_MATRIX) {
355     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
356     PetscCall(MatSetSizes(B, m, n, m, n));
357     PetscCall(MatSetType(B, MATSEQSBAIJ));
358     PetscCall(MatSeqSBAIJSetPreallocation(B, bs, 0, browlengths));
359   } else B = *newmat;
360 
361   b  = (Mat_SeqSBAIJ *)B->data;
362   bi = b->i;
363   bj = b->j;
364   bv = b->a;
365 
366   bi[0] = 0;
367   for (PetscInt i = 0; i < mbs; i++) {
368     aj = a->j + adiag[i];
369     av = a->a + (adiag[i]) * bs2;
370     for (PetscInt j = 0; j < browlengths[i]; j++) {
371       *bj = *aj;
372       bj++;
373       aj++;
374       for (k = 0; k < bs2; k++) {
375         *bv = *av;
376         bv++;
377         av++;
378       }
379     }
380     bi[i + 1]  = bi[i] + browlengths[i];
381     b->ilen[i] = browlengths[i];
382   }
383   PetscCall(PetscFree(browlengths));
384   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
385   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
386 
387   if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &B));
388   else *newmat = B;
389   PetscFunctionReturn(PETSC_SUCCESS);
390 }
391