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