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