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