xref: /petsc/src/mat/impls/sbaij/seq/aijsbaij.c (revision 43cdf1ebbcbfa05bee08e48007ef1bae3f20f4e9)
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(0);
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   PetscInt    m, n, bs = PetscAbs(A->rmap->bs);
130   PetscInt   *ii;
131   Mat_SeqAIJ *Aa = (Mat_SeqAIJ *)A->data;
132 
133   PetscFunctionBegin;
134   PetscCall(MatGetSize(A, &m, &n));
135   PetscCall(PetscCalloc1(m / bs, nnz));
136   PetscCall(PetscMalloc1(bs, &ii));
137 
138   /* loop over all potential blocks in the matrix to see if the potential block has a nonzero */
139   for (PetscInt i = 0; i < m / bs; i++) {
140     for (PetscInt k = 0; k < bs; k++) ii[k] = Aa->i[i * bs + k];
141     for (PetscInt j = 0; j < n / bs; j++) {
142       for (PetscInt k = 0; k < bs; k++) {
143         for (; ii[k] < Aa->i[i * bs + k + 1] && Aa->j[ii[k]] < (j + 1) * bs; ii[k]++) {
144           if (j >= i && Aa->j[ii[k]] >= j * bs) {
145             /* found a nonzero in the potential block so allocate for that block and jump to check the next potential block */
146             (*nnz)[i]++;
147             goto theend;
148           }
149         }
150       }
151     theend:;
152     }
153   }
154   PetscCall(PetscFree(ii));
155   PetscFunctionReturn(0);
156 }
157 
158 PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqSBAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
159 {
160   Mat           B;
161   Mat_SeqAIJ   *a = (Mat_SeqAIJ *)A->data;
162   Mat_SeqSBAIJ *b;
163   PetscInt     *ai = a->i, *aj, m = A->rmap->N, n = A->cmap->N, i, j, *bi, *bj, *rowlengths, bs = PetscAbs(A->rmap->bs);
164   MatScalar    *av, *bv;
165   PetscBool     miss = PETSC_FALSE;
166 
167   PetscFunctionBegin;
168 #if !defined(PETSC_USE_COMPLEX)
169   PetscCheck(A->symmetric == PETSC_BOOL3_TRUE, PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Matrix must be symmetric. Call MatSetOption(mat,MAT_SYMMETRIC,PETSC_TRUE)");
170 #else
171   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)");
172 #endif
173   PetscCheck(n == m, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Matrix must be square");
174 
175   if (bs == 1) {
176     PetscCall(PetscMalloc1(m, &rowlengths));
177     for (i = 0; i < m; i++) {
178       if (a->diag[i] == ai[i + 1]) {             /* missing diagonal */
179         rowlengths[i] = (ai[i + 1] - ai[i]) + 1; /* allocate some extra space */
180         miss          = PETSC_TRUE;
181       } else {
182         rowlengths[i] = (ai[i + 1] - a->diag[i]);
183       }
184     }
185   } else PetscCall(MatConvert_SeqAIJ_SeqSBAIJ_Preallocate(A, &rowlengths));
186   if (reuse != MAT_REUSE_MATRIX) {
187     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
188     PetscCall(MatSetSizes(B, m, n, m, n));
189     PetscCall(MatSetType(B, MATSEQSBAIJ));
190     PetscCall(MatSeqSBAIJSetPreallocation(B, bs, 0, rowlengths));
191   } else B = *newmat;
192 
193   if (bs == 1 && !miss) {
194     b  = (Mat_SeqSBAIJ *)(B->data);
195     bi = b->i;
196     bj = b->j;
197     bv = b->a;
198 
199     bi[0] = 0;
200     for (i = 0; i < m; i++) {
201       aj = a->j + a->diag[i];
202       av = a->a + a->diag[i];
203       for (j = 0; j < rowlengths[i]; j++) {
204         *bj = *aj;
205         bj++;
206         aj++;
207         *bv = *av;
208         bv++;
209         av++;
210       }
211       bi[i + 1]  = bi[i] + rowlengths[i];
212       b->ilen[i] = rowlengths[i];
213     }
214     PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
215     PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
216   } else {
217     PetscCall(MatSetOption(B, MAT_IGNORE_LOWER_TRIANGULAR, PETSC_TRUE));
218     /* reuse may not be equal to MAT_REUSE_MATRIX, but the basic converter will reallocate or replace newmat if this value is not used */
219     /* if reuse is equal to MAT_INITIAL_MATRIX, it has been appropriately preallocated before                                          */
220     /*                      MAT_INPLACE_MATRIX, it will be replaced with MatHeaderReplace below                                        */
221     PetscCall(MatConvert_Basic(A, newtype, MAT_REUSE_MATRIX, &B));
222   }
223   PetscCall(PetscFree(rowlengths));
224   if (reuse == MAT_INPLACE_MATRIX) {
225     PetscCall(MatHeaderReplace(A, &B));
226   } else *newmat = B;
227   PetscFunctionReturn(0);
228 }
229 
230 PETSC_INTERN PetscErrorCode MatConvert_SeqSBAIJ_SeqBAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
231 {
232   Mat           B;
233   Mat_SeqSBAIJ *a = (Mat_SeqSBAIJ *)A->data;
234   Mat_SeqBAIJ  *b;
235   PetscInt     *ai = a->i, *aj = a->j, m = A->rmap->N, n = A->cmap->n, i, k, *bi, *bj, *browlengths, nz, *browstart, itmp;
236   PetscInt      bs = A->rmap->bs, bs2 = bs * bs, mbs = m / bs, col, row;
237   MatScalar    *av, *bv;
238 #if defined(PETSC_USE_COMPLEX)
239   const int aconj = A->hermitian == PETSC_BOOL3_TRUE ? 1 : 0;
240 #else
241   const int aconj = 0;
242 #endif
243 
244   PetscFunctionBegin;
245   /* compute browlengths of newmat */
246   PetscCall(PetscMalloc2(mbs, &browlengths, mbs, &browstart));
247   for (i = 0; i < mbs; i++) browlengths[i] = 0;
248   for (i = 0; i < mbs; i++) {
249     nz = ai[i + 1] - ai[i];
250     aj++;                      /* skip diagonal */
251     for (k = 1; k < nz; k++) { /* no. of lower triangular blocks */
252       browlengths[*aj]++;
253       aj++;
254     }
255     browlengths[i] += nz; /* no. of upper triangular blocks */
256   }
257 
258   if (reuse != MAT_REUSE_MATRIX) {
259     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
260     PetscCall(MatSetSizes(B, m, n, m, n));
261     PetscCall(MatSetType(B, MATSEQBAIJ));
262     PetscCall(MatSeqBAIJSetPreallocation(B, bs, 0, browlengths));
263   } else B = *newmat;
264 
265   b  = (Mat_SeqBAIJ *)(B->data);
266   bi = b->i;
267   bj = b->j;
268   bv = b->a;
269 
270   /* set b->i */
271   bi[0] = 0;
272   for (i = 0; i < mbs; i++) {
273     b->ilen[i]   = browlengths[i];
274     bi[i + 1]    = bi[i] + browlengths[i];
275     browstart[i] = bi[i];
276   }
277   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);
278 
279   /* set b->j and b->a */
280   aj = a->j;
281   av = a->a;
282   for (i = 0; i < mbs; i++) {
283     /* diagonal block */
284     *(bj + browstart[i]) = *aj;
285     aj++;
286 
287     itmp = bs2 * browstart[i];
288     for (k = 0; k < bs2; k++) {
289       *(bv + itmp + k) = *av;
290       av++;
291     }
292     browstart[i]++;
293 
294     nz = ai[i + 1] - ai[i] - 1;
295     while (nz--) {
296       /* lower triangular blocks - transpose blocks of A */
297       *(bj + browstart[*aj]) = i; /* block col index */
298 
299       itmp = bs2 * browstart[*aj]; /* row index */
300       for (col = 0; col < bs; col++) {
301         k = col;
302         for (row = 0; row < bs; row++) {
303           bv[itmp + col * bs + row] = aconj ? PetscConj(av[k]) : av[k];
304           k += bs;
305         }
306       }
307       browstart[*aj]++;
308 
309       /* upper triangular blocks */
310       *(bj + browstart[i]) = *aj;
311       aj++;
312 
313       itmp = bs2 * browstart[i];
314       for (k = 0; k < bs2; k++) bv[itmp + k] = av[k];
315       av += bs2;
316       browstart[i]++;
317     }
318   }
319   PetscCall(PetscFree2(browlengths, browstart));
320   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
321   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
322 
323   if (reuse == MAT_INPLACE_MATRIX) {
324     PetscCall(MatHeaderReplace(A, &B));
325   } else *newmat = B;
326   PetscFunctionReturn(0);
327 }
328 
329 PETSC_INTERN PetscErrorCode MatConvert_SeqBAIJ_SeqSBAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
330 {
331   Mat           B;
332   Mat_SeqBAIJ  *a = (Mat_SeqBAIJ *)A->data;
333   Mat_SeqSBAIJ *b;
334   PetscInt     *ai = a->i, *aj, m = A->rmap->N, n = A->cmap->n, i, j, k, *bi, *bj, *browlengths;
335   PetscInt      bs = A->rmap->bs, bs2 = bs * bs, mbs = m / bs, dd;
336   MatScalar    *av, *bv;
337   PetscBool     flg;
338 
339   PetscFunctionBegin;
340   PetscCheck(A->symmetric, PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "Matrix must be symmetric. Call MatSetOption(mat,MAT_SYMMETRIC,PETSC_TRUE)");
341   PetscCheck(n == m, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Matrix must be square");
342   PetscCall(MatMissingDiagonal_SeqBAIJ(A, &flg, &dd)); /* check for missing diagonals, then mark diag */
343   PetscCheck(!flg, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix is missing diagonal %" PetscInt_FMT, dd);
344 
345   PetscCall(PetscMalloc1(mbs, &browlengths));
346   for (i = 0; i < mbs; i++) browlengths[i] = ai[i + 1] - a->diag[i];
347 
348   if (reuse != MAT_REUSE_MATRIX) {
349     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
350     PetscCall(MatSetSizes(B, m, n, m, n));
351     PetscCall(MatSetType(B, MATSEQSBAIJ));
352     PetscCall(MatSeqSBAIJSetPreallocation(B, bs, 0, browlengths));
353   } else B = *newmat;
354 
355   b  = (Mat_SeqSBAIJ *)(B->data);
356   bi = b->i;
357   bj = b->j;
358   bv = b->a;
359 
360   bi[0] = 0;
361   for (i = 0; i < mbs; i++) {
362     aj = a->j + a->diag[i];
363     av = a->a + (a->diag[i]) * bs2;
364     for (j = 0; j < browlengths[i]; j++) {
365       *bj = *aj;
366       bj++;
367       aj++;
368       for (k = 0; k < bs2; k++) {
369         *bv = *av;
370         bv++;
371         av++;
372       }
373     }
374     bi[i + 1]  = bi[i] + browlengths[i];
375     b->ilen[i] = browlengths[i];
376   }
377   PetscCall(PetscFree(browlengths));
378   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
379   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
380 
381   if (reuse == MAT_INPLACE_MATRIX) {
382     PetscCall(MatHeaderReplace(A, &B));
383   } else *newmat = B;
384   PetscFunctionReturn(0);
385 }
386