xref: /petsc/src/mat/impls/sbaij/seq/aijsbaij.c (revision bcee047adeeb73090d7e36cc71e39fc287cdbb97)
1 
2 #include <../src/mat/impls/aij/seq/aij.h>
3 #include <../src/mat/impls/baij/seq/baij.h>
4 #include <../src/mat/impls/sbaij/seq/sbaij.h>
5 
6 PETSC_INTERN PetscErrorCode MatConvert_SeqSBAIJ_SeqAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
7 {
8   Mat           B;
9   Mat_SeqSBAIJ *a = (Mat_SeqSBAIJ *)A->data;
10   Mat_SeqAIJ   *b;
11   PetscInt     *ai = a->i, *aj = a->j, m = A->rmap->N, n = A->cmap->n, i, j, k, *bi, *bj, *rowlengths, nz, *rowstart, itmp;
12   PetscInt      bs = A->rmap->bs, bs2 = bs * bs, mbs = A->rmap->N / bs, diagcnt = 0;
13   MatScalar    *av, *bv;
14 #if defined(PETSC_USE_COMPLEX)
15   const int aconj = A->hermitian == PETSC_BOOL3_TRUE ? 1 : 0;
16 #else
17   const int aconj = 0;
18 #endif
19 
20   PetscFunctionBegin;
21   /* compute rowlengths of newmat */
22   PetscCall(PetscMalloc2(m, &rowlengths, m + 1, &rowstart));
23 
24   for (i = 0; i < mbs; i++) rowlengths[i * bs] = 0;
25   k = 0;
26   for (i = 0; i < mbs; i++) {
27     nz = ai[i + 1] - ai[i];
28     if (nz) {
29       rowlengths[k] += nz; /* no. of upper triangular blocks */
30       if (*aj == i) {
31         aj++;
32         diagcnt++;
33         nz--;
34       }                          /* skip diagonal */
35       for (j = 0; j < nz; j++) { /* no. of lower triangular blocks */
36         rowlengths[(*aj) * bs]++;
37         aj++;
38       }
39     }
40     rowlengths[k] *= bs;
41     for (j = 1; j < bs; j++) rowlengths[k + j] = rowlengths[k];
42     k += bs;
43   }
44 
45   if (reuse != MAT_REUSE_MATRIX) {
46     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
47     PetscCall(MatSetSizes(B, m, n, m, n));
48     PetscCall(MatSetType(B, MATSEQAIJ));
49     PetscCall(MatSeqAIJSetPreallocation(B, 0, rowlengths));
50     PetscCall(MatSetBlockSize(B, A->rmap->bs));
51   } else B = *newmat;
52 
53   b  = (Mat_SeqAIJ *)(B->data);
54   bi = b->i;
55   bj = b->j;
56   bv = b->a;
57 
58   /* set b->i */
59   bi[0]       = 0;
60   rowstart[0] = 0;
61   for (i = 0; i < mbs; i++) {
62     for (j = 0; j < bs; j++) {
63       b->ilen[i * bs + j]      = rowlengths[i * bs];
64       rowstart[i * bs + j + 1] = rowstart[i * bs + j] + rowlengths[i * bs];
65     }
66     bi[i + 1] = bi[i] + rowlengths[i * bs] / bs;
67   }
68   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);
69 
70   /* set b->j and b->a */
71   aj = a->j;
72   av = a->a;
73   for (i = 0; i < mbs; i++) {
74     nz = ai[i + 1] - ai[i];
75     /* diagonal block */
76     if (nz && *aj == i) {
77       nz--;
78       for (j = 0; j < bs; j++) { /* row i*bs+j */
79         itmp = i * bs + j;
80         for (k = 0; k < bs; k++) { /* col i*bs+k */
81           *(bj + rowstart[itmp]) = (*aj) * bs + k;
82           *(bv + rowstart[itmp]) = *(av + k * bs + j);
83           rowstart[itmp]++;
84         }
85       }
86       aj++;
87       av += bs2;
88     }
89 
90     while (nz--) {
91       /* lower triangular blocks */
92       for (j = 0; j < bs; j++) { /* row (*aj)*bs+j */
93         itmp = (*aj) * bs + j;
94         for (k = 0; k < bs; k++) { /* col i*bs+k */
95           *(bj + rowstart[itmp]) = i * bs + k;
96           *(bv + rowstart[itmp]) = aconj ? PetscConj(*(av + j * bs + k)) : *(av + j * bs + k);
97           rowstart[itmp]++;
98         }
99       }
100       /* upper triangular blocks */
101       for (j = 0; j < bs; j++) { /* row i*bs+j */
102         itmp = i * bs + j;
103         for (k = 0; k < bs; k++) { /* col (*aj)*bs+k */
104           *(bj + rowstart[itmp]) = (*aj) * bs + k;
105           *(bv + rowstart[itmp]) = *(av + k * bs + j);
106           rowstart[itmp]++;
107         }
108       }
109       aj++;
110       av += bs2;
111     }
112   }
113   PetscCall(PetscFree2(rowlengths, rowstart));
114   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
115   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
116 
117   if (reuse == MAT_INPLACE_MATRIX) {
118     PetscCall(MatHeaderReplace(A, &B));
119   } else {
120     *newmat = B;
121   }
122   PetscFunctionReturn(PETSC_SUCCESS);
123 }
124 
125 #include <../src/mat/impls/aij/seq/aij.h>
126 
127 PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqSBAIJ_Preallocate(Mat A, PetscInt **nnz)
128 {
129   Mat_SeqAIJ     *Aa = (Mat_SeqAIJ *)A->data;
130   PetscInt        m, n, bs = PetscAbs(A->rmap->bs);
131   const PetscInt *ai = Aa->i, *aj = Aa->j;
132 
133   PetscFunctionBegin;
134   PetscCall(MatGetSize(A, &m, &n));
135 
136   if (bs == 1) {
137     const PetscInt *adiag = Aa->diag;
138 
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 = PetscAbs(A->rmap->bs);
173   MatScalar    *av, *bv;
174   PetscBool     miss = PETSC_FALSE;
175 
176   PetscFunctionBegin;
177 #if !defined(PETSC_USE_COMPLEX)
178   PetscCheck(A->symmetric == PETSC_BOOL3_TRUE, PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Matrix must be symmetric. Call MatSetOption(mat,MAT_SYMMETRIC,PETSC_TRUE)");
179 #else
180   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)");
181 #endif
182   PetscCheck(n == m, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Matrix must be square");
183 
184   if (bs == 1) {
185     PetscCall(PetscMalloc1(m, &rowlengths));
186     for (i = 0; i < m; i++) {
187       if (a->diag[i] == ai[i + 1]) {             /* missing diagonal */
188         rowlengths[i] = (ai[i + 1] - ai[i]) + 1; /* allocate some extra space */
189         miss          = PETSC_TRUE;
190       } else {
191         rowlengths[i] = (ai[i + 1] - a->diag[i]);
192       }
193     }
194   } else PetscCall(MatConvert_SeqAIJ_SeqSBAIJ_Preallocate(A, &rowlengths));
195   if (reuse != MAT_REUSE_MATRIX) {
196     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
197     PetscCall(MatSetSizes(B, m, n, m, n));
198     PetscCall(MatSetType(B, MATSEQSBAIJ));
199     PetscCall(MatSeqSBAIJSetPreallocation(B, bs, 0, rowlengths));
200   } else B = *newmat;
201 
202   if (bs == 1 && !miss) {
203     b  = (Mat_SeqSBAIJ *)(B->data);
204     bi = b->i;
205     bj = b->j;
206     bv = b->a;
207 
208     bi[0] = 0;
209     for (i = 0; i < m; i++) {
210       aj = a->j + a->diag[i];
211       av = a->a + a->diag[i];
212       for (j = 0; j < rowlengths[i]; j++) {
213         *bj = *aj;
214         bj++;
215         aj++;
216         *bv = *av;
217         bv++;
218         av++;
219       }
220       bi[i + 1]  = bi[i] + rowlengths[i];
221       b->ilen[i] = rowlengths[i];
222     }
223     PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
224     PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
225   } else {
226     PetscCall(MatSetOption(B, MAT_IGNORE_LOWER_TRIANGULAR, PETSC_TRUE));
227     /* reuse may not be equal to MAT_REUSE_MATRIX, but the basic converter will reallocate or replace newmat if this value is not used */
228     /* if reuse is equal to MAT_INITIAL_MATRIX, it has been appropriately preallocated before                                          */
229     /*                      MAT_INPLACE_MATRIX, it will be replaced with MatHeaderReplace below                                        */
230     PetscCall(MatConvert_Basic(A, newtype, MAT_REUSE_MATRIX, &B));
231   }
232   PetscCall(PetscFree(rowlengths));
233   if (reuse == MAT_INPLACE_MATRIX) {
234     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, i, k, *bi, *bj, *browlengths, nz, *browstart, itmp;
245   PetscInt      bs = A->rmap->bs, bs2 = bs * bs, mbs = m / bs, col, row;
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 (i = 0; i < mbs; i++) browlengths[i] = 0;
257   for (i = 0; i < mbs; i++) {
258     nz = ai[i + 1] - ai[i];
259     aj++;                      /* skip diagonal */
260     for (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 (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 (i = 0; i < mbs; i++) {
292     /* diagonal block */
293     *(bj + browstart[i]) = *aj;
294     aj++;
295 
296     itmp = bs2 * browstart[i];
297     for (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 (col = 0; col < bs; col++) {
310         k = col;
311         for (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 (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) {
333     PetscCall(MatHeaderReplace(A, &B));
334   } else *newmat = B;
335   PetscFunctionReturn(PETSC_SUCCESS);
336 }
337 
338 PETSC_INTERN PetscErrorCode MatConvert_SeqBAIJ_SeqSBAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
339 {
340   Mat           B;
341   Mat_SeqBAIJ  *a = (Mat_SeqBAIJ *)A->data;
342   Mat_SeqSBAIJ *b;
343   PetscInt     *ai = a->i, *aj, m = A->rmap->N, n = A->cmap->n, i, j, k, *bi, *bj, *browlengths;
344   PetscInt      bs = A->rmap->bs, bs2 = bs * bs, mbs = m / bs, dd;
345   MatScalar    *av, *bv;
346   PetscBool     flg;
347 
348   PetscFunctionBegin;
349   PetscCheck(A->symmetric, PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Matrix must be symmetric. Call MatSetOption(mat,MAT_SYMMETRIC,PETSC_TRUE)");
350   PetscCheck(n == m, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Matrix must be square");
351   PetscCall(MatMissingDiagonal_SeqBAIJ(A, &flg, &dd)); /* check for missing diagonals, then mark diag */
352   PetscCheck(!flg, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix is missing diagonal %" PetscInt_FMT, dd);
353 
354   PetscCall(PetscMalloc1(mbs, &browlengths));
355   for (i = 0; i < mbs; i++) browlengths[i] = ai[i + 1] - a->diag[i];
356 
357   if (reuse != MAT_REUSE_MATRIX) {
358     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
359     PetscCall(MatSetSizes(B, m, n, m, n));
360     PetscCall(MatSetType(B, MATSEQSBAIJ));
361     PetscCall(MatSeqSBAIJSetPreallocation(B, bs, 0, browlengths));
362   } else B = *newmat;
363 
364   b  = (Mat_SeqSBAIJ *)(B->data);
365   bi = b->i;
366   bj = b->j;
367   bv = b->a;
368 
369   bi[0] = 0;
370   for (i = 0; i < mbs; i++) {
371     aj = a->j + a->diag[i];
372     av = a->a + (a->diag[i]) * bs2;
373     for (j = 0; j < browlengths[i]; j++) {
374       *bj = *aj;
375       bj++;
376       aj++;
377       for (k = 0; k < bs2; k++) {
378         *bv = *av;
379         bv++;
380         av++;
381       }
382     }
383     bi[i + 1]  = bi[i] + browlengths[i];
384     b->ilen[i] = browlengths[i];
385   }
386   PetscCall(PetscFree(browlengths));
387   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
388   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
389 
390   if (reuse == MAT_INPLACE_MATRIX) {
391     PetscCall(MatHeaderReplace(A, &B));
392   } else *newmat = B;
393   PetscFunctionReturn(PETSC_SUCCESS);
394 }
395