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: 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: }
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 */
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: }
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;
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: }
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;
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: }
92: /* Compose an IS with an ISLocalToGlobalMapping to map from IS source indices to global indices */
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;
100: PetscFunctionBegin;
103: PetscAssertPointer(cltog, 3);
104: PetscCall(PetscObjectTypeCompare((PetscObject)is, ISBLOCK, &isblock));
105: if (isblock) {
106: PetscInt lbs;
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: }
135: static PetscErrorCode ISL2GComposeBlock(IS is, ISLocalToGlobalMapping ltog, ISLocalToGlobalMapping *cltog)
136: {
137: const PetscInt *idx;
138: PetscInt m, *idxm, bs;
140: PetscFunctionBegin;
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: }
158: static PetscErrorCode MatDestroy_LocalRef(Mat B)
159: {
160: PetscFunctionBegin;
161: PetscCall(PetscFree(B->data));
162: PetscFunctionReturn(PETSC_SUCCESS);
163: }
165: /*@
166: MatCreateLocalRef - Gets a logical reference to a local submatrix, for use in assembly, that is to set values into the matrix
168: Not Collective
170: Input Parameters:
171: + A - full matrix, generally parallel
172: . isrow - Local index set for the rows
173: - iscol - Local index set for the columns
175: Output Parameter:
176: . newmat - new serial `Mat`
178: Level: developer
180: Notes:
181: Most will use `MatGetLocalSubMatrix()` which returns a real matrix corresponding to the local
182: block if it available, such as with matrix formats that store these blocks separately.
184: The new matrix forwards `MatSetValuesLocal()` and `MatSetValuesBlockedLocal()` to the global system.
185: In general, it does not define `MatMult()` or any other functions. Local submatrices can be nested.
187: .seealso: [](ch_matrices), `Mat`, `MATSUBMATRIX`, `MatCreateSubMatrixVirtual()`, `MatSetValuesLocal()`, `MatSetValuesBlockedLocal()`, `MatGetLocalSubMatrix()`, `MatCreateSubMatrix()`
188: @*/
189: PetscErrorCode MatCreateLocalRef(Mat A, IS isrow, IS iscol, Mat *newmat)
190: {
191: Mat_LocalRef *lr;
192: Mat B;
193: PetscInt m, n;
194: PetscBool islr;
196: PetscFunctionBegin;
200: PetscAssertPointer(newmat, 4);
201: PetscCheck(A->rmap->mapping, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Matrix must have local to global mapping provided before this call");
202: *newmat = NULL;
204: PetscCall(MatCreate(PETSC_COMM_SELF, &B));
205: PetscCall(ISGetLocalSize(isrow, &m));
206: PetscCall(ISGetLocalSize(iscol, &n));
207: PetscCall(MatSetSizes(B, m, n, m, n));
208: PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATLOCALREF));
209: PetscCall(MatSetUp(B));
211: B->ops->destroy = MatDestroy_LocalRef;
213: PetscCall(PetscNew(&lr));
214: B->data = (void *)lr;
216: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATLOCALREF, &islr));
217: if (islr) {
218: Mat_LocalRef *alr = (Mat_LocalRef *)A->data;
219: lr->Top = alr->Top;
220: } else {
221: /* This does not increase the reference count because MatLocalRef is not allowed to live longer than its parent */
222: lr->Top = A;
223: }
224: {
225: ISLocalToGlobalMapping rltog, cltog;
226: PetscInt arbs, acbs, rbs, cbs;
228: /* We will translate directly to global indices for the top level */
229: lr->SetValues = MatSetValues;
230: lr->SetValuesBlocked = MatSetValuesBlocked;
232: B->ops->setvalueslocal = MatSetValuesLocal_LocalRef_Scalar;
234: PetscCall(ISL2GCompose(isrow, A->rmap->mapping, &rltog));
235: if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) {
236: PetscCall(PetscObjectReference((PetscObject)rltog));
237: cltog = rltog;
238: } else {
239: PetscCall(ISL2GCompose(iscol, A->cmap->mapping, &cltog));
240: }
241: /* Remember if the ISes we used to pull out the submatrix are of type ISBLOCK. Will be used later in
242: * MatSetValuesLocal_LocalRef_Scalar. */
243: PetscCall(PetscObjectTypeCompare((PetscObject)isrow, ISBLOCK, &lr->rowisblock));
244: PetscCall(PetscObjectTypeCompare((PetscObject)iscol, ISBLOCK, &lr->colisblock));
245: PetscCall(MatSetLocalToGlobalMapping(B, rltog, cltog));
246: PetscCall(ISLocalToGlobalMappingDestroy(&rltog));
247: PetscCall(ISLocalToGlobalMappingDestroy(&cltog));
249: PetscCall(MatGetBlockSizes(A, &arbs, &acbs));
250: PetscCall(ISGetBlockSize(isrow, &rbs));
251: PetscCall(ISGetBlockSize(iscol, &cbs));
252: /* Always support block interface insertion on submatrix */
253: PetscCall(PetscLayoutSetBlockSize(B->rmap, rbs));
254: PetscCall(PetscLayoutSetBlockSize(B->cmap, cbs));
255: if (arbs != rbs || acbs != cbs || (arbs == 1 && acbs == 1)) {
256: /* Top-level matrix has different block size, so we have to call its scalar insertion interface */
257: B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Scalar;
258: } else {
259: /* Block sizes match so we can forward values to the top level using the block interface */
260: B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Block;
262: PetscCall(ISL2GComposeBlock(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(ISL2GComposeBlock(iscol, A->cmap->mapping, &cltog));
268: }
269: PetscCall(MatSetLocalToGlobalMapping(B, rltog, cltog));
270: PetscCall(ISLocalToGlobalMappingDestroy(&rltog));
271: PetscCall(ISLocalToGlobalMappingDestroy(&cltog));
272: }
273: }
274: *newmat = B;
275: PetscFunctionReturn(PETSC_SUCCESS);
276: }