xref: /petsc/src/mat/impls/aij/seq/symtranspose.c (revision d52a580b706c59ca78066c1e38754e45b6b56e2b)
1 /*
2   Defines transpose routines for SeqAIJ matrices.
3 */
4 
5 #include <../src/mat/impls/aij/seq/aij.h>
6 
7 /*
8    The symbolic and full transpose versions share several similar code blocks but the macros to reuse the code would be confusing and ugly
9 */
10 PetscErrorCode MatTransposeSymbolic_SeqAIJ(Mat A, Mat *B)
11 {
12   PetscInt    i, j, anzj;
13   Mat         At;
14   Mat_SeqAIJ *a  = (Mat_SeqAIJ *)A->data, *at;
15   PetscInt    an = A->cmap->N, am = A->rmap->N;
16   PetscInt   *ati, *atj, *atfill, *ai = a->i, *aj = a->j;
17 
18   PetscFunctionBegin;
19   /* Allocate space for symbolic transpose info and work array */
20   PetscCall(PetscCalloc1(an + 1, &ati));
21   PetscCall(PetscMalloc1(ai[am], &atj));
22 
23   /* Walk through aj and count ## of non-zeros in each row of A^T. */
24   /* Note: offset by 1 for fast conversion into csr format. */
25   for (i = 0; i < ai[am]; i++) ati[aj[i] + 1] += 1;
26   /* Form ati for csr format of A^T. */
27   for (i = 0; i < an; i++) ati[i + 1] += ati[i];
28 
29   /* Copy ati into atfill so we have locations of the next free space in atj */
30   PetscCall(PetscMalloc1(an, &atfill));
31   PetscCall(PetscArraycpy(atfill, ati, an));
32 
33   /* Walk through A row-wise and mark nonzero entries of A^T. */
34   for (i = 0; i < am; i++) {
35     anzj = ai[i + 1] - ai[i];
36     for (j = 0; j < anzj; j++) {
37       atj[atfill[*aj]] = i;
38       atfill[*aj++] += 1;
39     }
40   }
41   PetscCall(PetscFree(atfill));
42 
43   PetscCall(MatCreateSeqAIJWithArrays(PetscObjectComm((PetscObject)A), an, am, ati, atj, NULL, &At));
44   PetscCall(MatSetBlockSizes(At, A->cmap->bs, A->rmap->bs));
45   PetscCall(MatSetType(At, ((PetscObject)A)->type_name));
46   at          = (Mat_SeqAIJ *)At->data;
47   at->free_a  = PETSC_FALSE;
48   at->free_ij = PETSC_TRUE;
49   at->nonew   = 0;
50   at->maxnz   = ati[an];
51   *B          = At;
52   PetscFunctionReturn(PETSC_SUCCESS);
53 }
54 
55 PetscErrorCode MatTranspose_SeqAIJ(Mat A, MatReuse reuse, Mat *B)
56 {
57   PetscInt         i, j, anzj;
58   Mat              At;
59   Mat_SeqAIJ      *a  = (Mat_SeqAIJ *)A->data, *at;
60   PetscInt         an = A->cmap->N, am = A->rmap->N;
61   PetscInt        *ati, *atj, *atfill, *ai = a->i, *aj = a->j;
62   MatScalar       *ata;
63   const MatScalar *aa, *av;
64   PetscContainer   rB;
65   MatParentState  *rb;
66   PetscBool        nonzerochange = PETSC_FALSE;
67 
68   PetscFunctionBegin;
69   if (reuse == MAT_REUSE_MATRIX) {
70     PetscCall(PetscObjectQuery((PetscObject)*B, "MatTransposeParent", (PetscObject *)&rB));
71     PetscCheck(rB, PetscObjectComm((PetscObject)*B), PETSC_ERR_ARG_WRONG, "Reuse matrix used was not generated from call to MatTranspose()");
72     PetscCall(PetscContainerGetPointer(rB, &rb));
73     if (rb->nonzerostate != A->nonzerostate) nonzerochange = PETSC_TRUE;
74   }
75 
76   PetscCall(MatSeqAIJGetArrayRead(A, &av));
77   aa = av;
78   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX || nonzerochange) {
79     /* Allocate space for symbolic transpose info and work array */
80     PetscCall(PetscCalloc1(an + 1, &ati));
81     PetscCall(PetscMalloc1(ai[am], &atj));
82     /* Walk through aj and count ## of non-zeros in each row of A^T. */
83     /* Note: offset by 1 for fast conversion into csr format. */
84     for (i = 0; i < ai[am]; i++) ati[aj[i] + 1] += 1;
85     /* Form ati for csr format of A^T. */
86     for (i = 0; i < an; i++) ati[i + 1] += ati[i];
87     PetscCall(PetscMalloc1(ai[am], &ata));
88   } else {
89     Mat_SeqAIJ *sub_B = (Mat_SeqAIJ *)(*B)->data;
90     ati               = sub_B->i;
91     atj               = sub_B->j;
92     ata               = sub_B->a;
93     At                = *B;
94   }
95 
96   /* Copy ati into atfill so we have locations of the next free space in atj */
97   PetscCall(PetscMalloc1(an, &atfill));
98   PetscCall(PetscArraycpy(atfill, ati, an));
99 
100   /* Walk through A row-wise and mark nonzero entries of A^T. */
101   if (aa) {
102     for (i = 0; i < am; i++) {
103       anzj = ai[i + 1] - ai[i];
104       for (j = 0; j < anzj; j++) {
105         atj[atfill[*aj]] = i;
106         ata[atfill[*aj]] = *aa++;
107         atfill[*aj++] += 1;
108       }
109     }
110   } else {
111     for (i = 0; i < am; i++) {
112       anzj = ai[i + 1] - ai[i];
113       for (j = 0; j < anzj; j++) {
114         atj[atfill[*aj]] = i;
115         atfill[*aj++] += 1;
116       }
117     }
118   }
119   PetscCall(PetscFree(atfill));
120   PetscCall(MatSeqAIJRestoreArrayRead(A, &av));
121   if (reuse == MAT_REUSE_MATRIX) PetscCall(PetscObjectStateIncrease((PetscObject)*B));
122 
123   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX || nonzerochange) {
124     PetscCall(MatCreateSeqAIJWithArrays(PetscObjectComm((PetscObject)A), an, am, ati, atj, ata, &At));
125     PetscCall(MatSetBlockSizes(At, A->cmap->bs, A->rmap->bs));
126     PetscCall(MatSetType(At, ((PetscObject)A)->type_name));
127     at          = (Mat_SeqAIJ *)At->data;
128     at->free_a  = PETSC_TRUE;
129     at->free_ij = PETSC_TRUE;
130     at->nonew   = 0;
131     at->maxnz   = ati[an];
132   }
133 
134   if (reuse == MAT_INITIAL_MATRIX || (reuse == MAT_REUSE_MATRIX && !nonzerochange)) {
135     *B = At;
136   } else if (nonzerochange) {
137     PetscCall(MatHeaderMerge(*B, &At));
138     PetscCall(MatTransposeSetPrecursor(A, *B));
139   } else if (reuse == MAT_INPLACE_MATRIX) {
140     PetscCall(MatHeaderMerge(A, &At));
141   }
142   PetscFunctionReturn(PETSC_SUCCESS);
143 }
144 
145 /*
146    Get symbolic matrix nonzero structure of a submatrix of A, A[rstart:rend,:],
147 */
148 PetscErrorCode MatGetSymbolicTransposeReduced_SeqAIJ(Mat A, PetscInt rstart, PetscInt rend, PetscInt *Ati[], PetscInt *Atj[])
149 {
150   PetscInt    i, j, anzj;
151   Mat_SeqAIJ *a  = (Mat_SeqAIJ *)A->data;
152   PetscInt    an = A->cmap->N;
153   PetscInt   *ati, *atj, *atfill, *ai = a->i, *aj = a->j, am = ai[rend] - ai[rstart];
154 
155   PetscFunctionBegin;
156   PetscCall(PetscLogEventBegin(MAT_Getsymtransreduced, A, 0, 0, 0));
157 
158   /* Allocate space for symbolic transpose info and work array */
159   PetscCall(PetscCalloc1(an + 1, &ati));
160   PetscCall(PetscMalloc1(am + 1, &atj));
161 
162   /* Walk through aj and count ## of non-zeros in each row of A^T. */
163   /* Note: offset by 1 for fast conversion into csr format. */
164   for (i = ai[rstart]; i < ai[rend]; i++) ati[aj[i] + 1] += 1;
165   /* Form ati for csr format of A^T. */
166   for (i = 0; i < an; i++) ati[i + 1] += ati[i];
167 
168   /* Copy ati into atfill so we have locations of the next free space in atj */
169   PetscCall(PetscMalloc1(an + 1, &atfill));
170   PetscCall(PetscArraycpy(atfill, ati, an));
171 
172   /* Walk through A row-wise and mark nonzero entries of A^T. */
173   aj = PetscSafePointerPlusOffset(aj, ai[rstart]);
174   for (i = rstart; i < rend; i++) {
175     anzj = ai[i + 1] - ai[i];
176     for (j = 0; j < anzj; j++) {
177       atj[atfill[*aj]] = i - rstart;
178       atfill[*aj++] += 1;
179     }
180   }
181   PetscCall(PetscFree(atfill));
182   *Ati = ati;
183   *Atj = atj;
184 
185   PetscCall(PetscLogEventEnd(MAT_Getsymtransreduced, A, 0, 0, 0));
186   PetscFunctionReturn(PETSC_SUCCESS);
187 }
188 
189 /*
190     Returns the i and j arrays for a symbolic transpose, this is used internally within SeqAIJ code when the full
191     symbolic matrix (which can be obtained with MatTransposeSymbolic() is not needed. MatRestoreSymbolicTranspose_SeqAIJ() should be used to free the arrays.
192 */
193 PetscErrorCode MatGetSymbolicTranspose_SeqAIJ(Mat A, PetscInt *Ati[], PetscInt *Atj[])
194 {
195   PetscFunctionBegin;
196   PetscCall(MatGetSymbolicTransposeReduced_SeqAIJ(A, 0, A->rmap->N, Ati, Atj));
197   PetscFunctionReturn(PETSC_SUCCESS);
198 }
199 
200 PetscErrorCode MatRestoreSymbolicTranspose_SeqAIJ(Mat A, PetscInt *ati[], PetscInt *atj[])
201 {
202   PetscFunctionBegin;
203   PetscCall(PetscFree(*ati));
204   PetscCall(PetscFree(*atj));
205   PetscFunctionReturn(PETSC_SUCCESS);
206 }
207