1 #include <petsc/private/matimpl.h> /*I "petscmat.h" I*/
2
3 typedef struct {
4 Mat Top;
5 PetscBool rowisblock;
6 PetscBool colisblock;
7 PetscErrorCode (*SetValues)(Mat, PetscInt, const PetscInt[], PetscInt, const PetscInt[], const PetscScalar[], InsertMode);
8 PetscErrorCode (*SetValuesBlocked)(Mat, PetscInt, const PetscInt[], PetscInt, const PetscInt[], const PetscScalar[], InsertMode);
9 } Mat_LocalRef;
10
11 /* These need to be macros because they use sizeof */
12 #define IndexSpaceGet(buf, nrow, ncol, irowm, icolm) \
13 do { \
14 if (nrow + ncol > (PetscInt)PETSC_STATIC_ARRAY_LENGTH(buf)) { \
15 PetscCall(PetscMalloc2(nrow, &irowm, ncol, &icolm)); \
16 } else { \
17 irowm = &buf[0]; \
18 icolm = &buf[nrow]; \
19 } \
20 } while (0)
21
22 #define IndexSpaceRestore(buf, nrow, ncol, irowm, icolm) \
23 do { \
24 if (nrow + ncol > (PetscInt)PETSC_STATIC_ARRAY_LENGTH(buf)) PetscCall(PetscFree2(irowm, icolm)); \
25 } while (0)
26
BlockIndicesExpand(PetscInt n,const PetscInt idx[],PetscInt bs,PetscInt idxm[])27 static void BlockIndicesExpand(PetscInt n, const PetscInt idx[], PetscInt bs, PetscInt idxm[])
28 {
29 PetscInt i, j;
30 for (i = 0; i < n; i++) {
31 for (j = 0; j < bs; j++) idxm[i * bs + j] = idx[i] * bs + j;
32 }
33 }
34
MatSetValuesBlockedLocal_LocalRef_Block(Mat A,PetscInt nrow,const PetscInt irow[],PetscInt ncol,const PetscInt icol[],const PetscScalar y[],InsertMode addv)35 static PetscErrorCode MatSetValuesBlockedLocal_LocalRef_Block(Mat A, PetscInt nrow, const PetscInt irow[], PetscInt ncol, const PetscInt icol[], const PetscScalar y[], InsertMode addv)
36 {
37 Mat_LocalRef *lr = (Mat_LocalRef *)A->data;
38 PetscInt buf[4096], *irowm = NULL, *icolm; /* suppress maybe-uninitialized warning */
39
40 PetscFunctionBegin;
41 if (!nrow || !ncol) PetscFunctionReturn(PETSC_SUCCESS);
42 IndexSpaceGet(buf, nrow, ncol, irowm, icolm);
43 PetscCall(ISLocalToGlobalMappingApplyBlock(A->rmap->mapping, nrow, irow, irowm));
44 PetscCall(ISLocalToGlobalMappingApplyBlock(A->cmap->mapping, ncol, icol, icolm));
45 PetscCall((*lr->SetValuesBlocked)(lr->Top, nrow, irowm, ncol, icolm, y, addv));
46 IndexSpaceRestore(buf, nrow, ncol, irowm, icolm);
47 PetscFunctionReturn(PETSC_SUCCESS);
48 }
49
MatSetValuesBlockedLocal_LocalRef_Scalar(Mat A,PetscInt nrow,const PetscInt irow[],PetscInt ncol,const PetscInt icol[],const PetscScalar y[],InsertMode addv)50 static PetscErrorCode MatSetValuesBlockedLocal_LocalRef_Scalar(Mat A, PetscInt nrow, const PetscInt irow[], PetscInt ncol, const PetscInt icol[], const PetscScalar y[], InsertMode addv)
51 {
52 Mat_LocalRef *lr = (Mat_LocalRef *)A->data;
53 PetscInt rbs, cbs, buf[4096], *irowm, *icolm;
54
55 PetscFunctionBegin;
56 PetscCall(MatGetBlockSizes(A, &rbs, &cbs));
57 IndexSpaceGet(buf, nrow * rbs, ncol * cbs, irowm, icolm);
58 BlockIndicesExpand(nrow, irow, rbs, irowm);
59 BlockIndicesExpand(ncol, icol, cbs, icolm);
60 PetscCall(ISLocalToGlobalMappingApplyBlock(A->rmap->mapping, nrow * rbs, irowm, irowm));
61 PetscCall(ISLocalToGlobalMappingApplyBlock(A->cmap->mapping, ncol * cbs, icolm, icolm));
62 PetscCall((*lr->SetValues)(lr->Top, nrow * rbs, irowm, ncol * cbs, icolm, y, addv));
63 IndexSpaceRestore(buf, nrow * rbs, ncol * cbs, irowm, icolm);
64 PetscFunctionReturn(PETSC_SUCCESS);
65 }
66
MatSetValuesLocal_LocalRef_Scalar(Mat A,PetscInt nrow,const PetscInt irow[],PetscInt ncol,const PetscInt icol[],const PetscScalar y[],InsertMode addv)67 static PetscErrorCode MatSetValuesLocal_LocalRef_Scalar(Mat A, PetscInt nrow, const PetscInt irow[], PetscInt ncol, const PetscInt icol[], const PetscScalar y[], InsertMode addv)
68 {
69 Mat_LocalRef *lr = (Mat_LocalRef *)A->data;
70 PetscInt buf[4096], *irowm, *icolm;
71
72 PetscFunctionBegin;
73 IndexSpaceGet(buf, nrow, ncol, irowm, icolm);
74 /* If the row IS defining this submatrix was an ISBLOCK, then the unblocked LGMapApply is the right one to use. If
75 * instead it was (say) an ISSTRIDE with a block size > 1, then we need to use LGMapApplyBlock */
76 if (lr->rowisblock) {
77 PetscCall(ISLocalToGlobalMappingApply(A->rmap->mapping, nrow, irow, irowm));
78 } else {
79 PetscCall(ISLocalToGlobalMappingApplyBlock(A->rmap->mapping, nrow, irow, irowm));
80 }
81 /* As above, but for the column IS. */
82 if (lr->colisblock) {
83 PetscCall(ISLocalToGlobalMappingApply(A->cmap->mapping, ncol, icol, icolm));
84 } else {
85 PetscCall(ISLocalToGlobalMappingApplyBlock(A->cmap->mapping, ncol, icol, icolm));
86 }
87 PetscCall((*lr->SetValues)(lr->Top, nrow, irowm, ncol, icolm, y, addv));
88 IndexSpaceRestore(buf, nrow, ncol, irowm, icolm);
89 PetscFunctionReturn(PETSC_SUCCESS);
90 }
91
92 /* Compose an IS with an ISLocalToGlobalMapping to map from IS source indices to global indices */
ISL2GCompose(IS is,ISLocalToGlobalMapping ltog,ISLocalToGlobalMapping * cltog)93 static PetscErrorCode ISL2GCompose(IS is, ISLocalToGlobalMapping ltog, ISLocalToGlobalMapping *cltog)
94 {
95 const PetscInt *idx;
96 PetscInt m, *idxm;
97 PetscInt bs;
98 PetscBool isblock;
99
100 PetscFunctionBegin;
101 PetscValidHeaderSpecific(is, IS_CLASSID, 1);
102 PetscValidHeaderSpecific(ltog, IS_LTOGM_CLASSID, 2);
103 PetscAssertPointer(cltog, 3);
104 PetscCall(PetscObjectTypeCompare((PetscObject)is, ISBLOCK, &isblock));
105 if (isblock) {
106 PetscInt lbs;
107
108 PetscCall(ISGetBlockSize(is, &bs));
109 PetscCall(ISLocalToGlobalMappingGetBlockSize(ltog, &lbs));
110 if (bs == lbs) {
111 PetscCall(ISGetLocalSize(is, &m));
112 m = m / bs;
113 PetscCall(ISBlockGetIndices(is, &idx));
114 PetscCall(PetscMalloc1(m, &idxm));
115 PetscCall(ISLocalToGlobalMappingApplyBlock(ltog, m, idx, idxm));
116 PetscCall(ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is), bs, m, idxm, PETSC_OWN_POINTER, cltog));
117 PetscCall(ISBlockRestoreIndices(is, &idx));
118 PetscFunctionReturn(PETSC_SUCCESS);
119 }
120 }
121 PetscCall(ISGetLocalSize(is, &m));
122 PetscCall(ISGetIndices(is, &idx));
123 PetscCall(ISGetBlockSize(is, &bs));
124 PetscCall(PetscMalloc1(m, &idxm));
125 if (ltog) {
126 PetscCall(ISLocalToGlobalMappingApply(ltog, m, idx, idxm));
127 } else {
128 PetscCall(PetscArraycpy(idxm, idx, m));
129 }
130 PetscCall(ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is), bs, m, idxm, PETSC_OWN_POINTER, cltog));
131 PetscCall(ISRestoreIndices(is, &idx));
132 PetscFunctionReturn(PETSC_SUCCESS);
133 }
134
ISL2GComposeBlock(IS is,ISLocalToGlobalMapping ltog,ISLocalToGlobalMapping * cltog)135 static PetscErrorCode ISL2GComposeBlock(IS is, ISLocalToGlobalMapping ltog, ISLocalToGlobalMapping *cltog)
136 {
137 const PetscInt *idx;
138 PetscInt m, *idxm, bs;
139
140 PetscFunctionBegin;
141 PetscValidHeaderSpecific(is, IS_CLASSID, 1);
142 PetscValidHeaderSpecific(ltog, IS_LTOGM_CLASSID, 2);
143 PetscAssertPointer(cltog, 3);
144 PetscCall(ISBlockGetLocalSize(is, &m));
145 PetscCall(ISBlockGetIndices(is, &idx));
146 PetscCall(ISLocalToGlobalMappingGetBlockSize(ltog, &bs));
147 PetscCall(PetscMalloc1(m, &idxm));
148 if (ltog) {
149 PetscCall(ISLocalToGlobalMappingApplyBlock(ltog, m, idx, idxm));
150 } else {
151 PetscCall(PetscArraycpy(idxm, idx, m));
152 }
153 PetscCall(ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is), bs, m, idxm, PETSC_OWN_POINTER, cltog));
154 PetscCall(ISBlockRestoreIndices(is, &idx));
155 PetscFunctionReturn(PETSC_SUCCESS);
156 }
157
MatZeroRowsLocal_LocalRef(Mat A,PetscInt n,const PetscInt rows[],PetscScalar diag,Vec x,Vec b)158 static PetscErrorCode MatZeroRowsLocal_LocalRef(Mat A, PetscInt n, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
159 {
160 PetscInt *rows_l;
161 Mat_LocalRef *lr = (Mat_LocalRef *)A->data;
162
163 PetscFunctionBegin;
164 PetscCall(PetscMalloc1(n, &rows_l));
165 PetscCall(ISLocalToGlobalMappingApply(A->rmap->mapping, n, rows, rows_l));
166 PetscCall(MatZeroRows(lr->Top, n, rows_l, diag, x, b));
167 PetscCall(PetscFree(rows_l));
168 PetscFunctionReturn(PETSC_SUCCESS);
169 }
170
MatZeroRowsColumnsLocal_LocalRef(Mat A,PetscInt n,const PetscInt rows[],PetscScalar diag,Vec x,Vec b)171 static PetscErrorCode MatZeroRowsColumnsLocal_LocalRef(Mat A, PetscInt n, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
172 {
173 PetscInt *rows_l;
174 Mat_LocalRef *lr = (Mat_LocalRef *)A->data;
175
176 PetscFunctionBegin;
177 PetscCall(PetscMalloc1(n, &rows_l));
178 PetscCall(ISLocalToGlobalMappingApply(A->rmap->mapping, n, rows, rows_l));
179 PetscCall(MatZeroRowsColumns(lr->Top, n, rows_l, diag, x, b));
180 PetscCall(PetscFree(rows_l));
181 PetscFunctionReturn(PETSC_SUCCESS);
182 }
183
MatDestroy_LocalRef(Mat B)184 static PetscErrorCode MatDestroy_LocalRef(Mat B)
185 {
186 PetscFunctionBegin;
187 PetscCall(PetscFree(B->data));
188 PetscFunctionReturn(PETSC_SUCCESS);
189 }
190
191 /*@
192 MatCreateLocalRef - Gets a logical reference to a local submatrix, for use in assembly, that is to set values into the matrix
193
194 Not Collective
195
196 Input Parameters:
197 + A - full matrix, generally parallel
198 . isrow - Local index set for the rows
199 - iscol - Local index set for the columns
200
201 Output Parameter:
202 . newmat - new serial `Mat`
203
204 Level: developer
205
206 Notes:
207 Most will use `MatGetLocalSubMatrix()` which returns a real matrix corresponding to the local
208 block if it available, such as with matrix formats that store these blocks separately.
209
210 The new matrix forwards `MatSetValuesLocal()` and `MatSetValuesBlockedLocal()` to the global system.
211 In general, it does not define `MatMult()` or any other functions. Local submatrices can be nested.
212
213 .seealso: [](ch_matrices), `Mat`, `MATSUBMATRIX`, `MatCreateSubMatrixVirtual()`, `MatSetValuesLocal()`, `MatSetValuesBlockedLocal()`, `MatGetLocalSubMatrix()`, `MatCreateSubMatrix()`
214 @*/
MatCreateLocalRef(Mat A,IS isrow,IS iscol,Mat * newmat)215 PetscErrorCode MatCreateLocalRef(Mat A, IS isrow, IS iscol, Mat *newmat)
216 {
217 Mat_LocalRef *lr;
218 Mat B;
219 PetscInt m, n;
220 PetscBool islr;
221
222 PetscFunctionBegin;
223 PetscValidHeaderSpecific(A, MAT_CLASSID, 1);
224 PetscValidHeaderSpecific(isrow, IS_CLASSID, 2);
225 PetscValidHeaderSpecific(iscol, IS_CLASSID, 3);
226 PetscAssertPointer(newmat, 4);
227 PetscCheck(A->rmap->mapping, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Matrix must have local to global mapping provided before this call");
228 *newmat = NULL;
229
230 PetscCall(MatCreate(PETSC_COMM_SELF, &B));
231 PetscCall(ISGetLocalSize(isrow, &m));
232 PetscCall(ISGetLocalSize(iscol, &n));
233 PetscCall(MatSetSizes(B, m, n, m, n));
234 PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATLOCALREF));
235 PetscCall(MatSetUp(B));
236
237 B->ops->destroy = MatDestroy_LocalRef;
238
239 PetscCall(PetscNew(&lr));
240 B->data = (void *)lr;
241
242 PetscCall(PetscObjectTypeCompare((PetscObject)A, MATLOCALREF, &islr));
243 if (islr) {
244 Mat_LocalRef *alr = (Mat_LocalRef *)A->data;
245 lr->Top = alr->Top;
246 } else {
247 /* This does not increase the reference count because MatLocalRef is not allowed to live longer than its parent */
248 lr->Top = A;
249 }
250 {
251 ISLocalToGlobalMapping rltog, cltog;
252 PetscInt arbs, acbs, rbs, cbs;
253
254 /* We will translate directly to global indices for the top level */
255 lr->SetValues = MatSetValues;
256 lr->SetValuesBlocked = MatSetValuesBlocked;
257
258 B->ops->setvalueslocal = MatSetValuesLocal_LocalRef_Scalar;
259 B->ops->zerorowslocal = MatZeroRowsLocal_LocalRef;
260 B->ops->zerorowscolumnslocal = MatZeroRowsColumnsLocal_LocalRef;
261
262 PetscCall(ISL2GCompose(isrow, A->rmap->mapping, &rltog));
263 if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) {
264 PetscCall(PetscObjectReference((PetscObject)rltog));
265 cltog = rltog;
266 } else {
267 PetscCall(ISL2GCompose(iscol, A->cmap->mapping, &cltog));
268 }
269 /* Remember if the ISes we used to pull out the submatrix are of type ISBLOCK. Will be used later in
270 * MatSetValuesLocal_LocalRef_Scalar. */
271 PetscCall(PetscObjectTypeCompare((PetscObject)isrow, ISBLOCK, &lr->rowisblock));
272 PetscCall(PetscObjectTypeCompare((PetscObject)iscol, ISBLOCK, &lr->colisblock));
273 PetscCall(MatSetLocalToGlobalMapping(B, rltog, cltog));
274 PetscCall(ISLocalToGlobalMappingDestroy(&rltog));
275 PetscCall(ISLocalToGlobalMappingDestroy(&cltog));
276
277 PetscCall(MatGetBlockSizes(A, &arbs, &acbs));
278 PetscCall(ISGetBlockSize(isrow, &rbs));
279 PetscCall(ISGetBlockSize(iscol, &cbs));
280 /* Always support block interface insertion on submatrix */
281 PetscCall(PetscLayoutSetBlockSize(B->rmap, rbs));
282 PetscCall(PetscLayoutSetBlockSize(B->cmap, cbs));
283 if (arbs != rbs || acbs != cbs || (arbs == 1 && acbs == 1)) {
284 /* Top-level matrix has different block size, so we have to call its scalar insertion interface */
285 B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Scalar;
286 } else {
287 /* Block sizes match so we can forward values to the top level using the block interface */
288 B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Block;
289
290 PetscCall(ISL2GComposeBlock(isrow, A->rmap->mapping, &rltog));
291 if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) {
292 PetscCall(PetscObjectReference((PetscObject)rltog));
293 cltog = rltog;
294 } else {
295 PetscCall(ISL2GComposeBlock(iscol, A->cmap->mapping, &cltog));
296 }
297 PetscCall(MatSetLocalToGlobalMapping(B, rltog, cltog));
298 PetscCall(ISLocalToGlobalMappingDestroy(&rltog));
299 PetscCall(ISLocalToGlobalMappingDestroy(&cltog));
300 }
301 }
302 *newmat = B;
303 PetscFunctionReturn(PETSC_SUCCESS);
304 }
305