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: }