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:       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)) 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:   if (!nrow || !ncol) return 0;
 42:   IndexSpaceGet(buf, nrow, ncol, irowm, icolm);
 43:   ISLocalToGlobalMappingApplyBlock(A->rmap->mapping, nrow, irow, irowm);
 44:   ISLocalToGlobalMappingApplyBlock(A->cmap->mapping, ncol, icol, icolm);
 45:   (*lr->SetValuesBlocked)(lr->Top, nrow, irowm, ncol, icolm, y, addv);
 46:   IndexSpaceRestore(buf, nrow, ncol, irowm, icolm);
 47:   return 0;
 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:   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:   ISLocalToGlobalMappingApplyBlock(A->rmap->mapping, nrow * rbs, irowm, irowm);
 60:   ISLocalToGlobalMappingApplyBlock(A->cmap->mapping, ncol * cbs, icolm, icolm);
 61:   (*lr->SetValues)(lr->Top, nrow * rbs, irowm, ncol * cbs, icolm, y, addv);
 62:   IndexSpaceRestore(buf, nrow * rbs, ncol * cbs, irowm, icolm);
 63:   return 0;
 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:   IndexSpaceGet(buf, nrow, ncol, irowm, icolm);
 72:   /* If the row IS defining this submatrix was an ISBLOCK, then the unblocked LGMapApply is the right one to use.  If
 73:    * instead it was (say) an ISSTRIDE with a block size > 1, then we need to use LGMapApplyBlock */
 74:   if (lr->rowisblock) {
 75:     ISLocalToGlobalMappingApply(A->rmap->mapping, nrow, irow, irowm);
 76:   } else {
 77:     ISLocalToGlobalMappingApplyBlock(A->rmap->mapping, nrow, irow, irowm);
 78:   }
 79:   /* As above, but for the column IS. */
 80:   if (lr->colisblock) {
 81:     ISLocalToGlobalMappingApply(A->cmap->mapping, ncol, icol, icolm);
 82:   } else {
 83:     ISLocalToGlobalMappingApplyBlock(A->cmap->mapping, ncol, icol, icolm);
 84:   }
 85:   (*lr->SetValues)(lr->Top, nrow, irowm, ncol, icolm, y, addv);
 86:   IndexSpaceRestore(buf, nrow, ncol, irowm, icolm);
 87:   return 0;
 88: }

 90: /* Compose an IS with an ISLocalToGlobalMapping to map from IS source indices to global indices */
 91: static PetscErrorCode ISL2GCompose(IS is, ISLocalToGlobalMapping ltog, ISLocalToGlobalMapping *cltog)
 92: {
 93:   const PetscInt *idx;
 94:   PetscInt        m, *idxm;
 95:   PetscInt        bs;
 96:   PetscBool       isblock;

101:   PetscObjectTypeCompare((PetscObject)is, ISBLOCK, &isblock);
102:   if (isblock) {
103:     PetscInt lbs;

105:     ISGetBlockSize(is, &bs);
106:     ISLocalToGlobalMappingGetBlockSize(ltog, &lbs);
107:     if (bs == lbs) {
108:       ISGetLocalSize(is, &m);
109:       m = m / bs;
110:       ISBlockGetIndices(is, &idx);
111:       PetscMalloc1(m, &idxm);
112:       ISLocalToGlobalMappingApplyBlock(ltog, m, idx, idxm);
113:       ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is), bs, m, idxm, PETSC_OWN_POINTER, cltog);
114:       ISBlockRestoreIndices(is, &idx);
115:       return 0;
116:     }
117:   }
118:   ISGetLocalSize(is, &m);
119:   ISGetIndices(is, &idx);
120:   ISGetBlockSize(is, &bs);
121:   PetscMalloc1(m, &idxm);
122:   if (ltog) {
123:     ISLocalToGlobalMappingApply(ltog, m, idx, idxm);
124:   } else {
125:     PetscArraycpy(idxm, idx, m);
126:   }
127:   ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is), bs, m, idxm, PETSC_OWN_POINTER, cltog);
128:   ISRestoreIndices(is, &idx);
129:   return 0;
130: }

132: static PetscErrorCode ISL2GComposeBlock(IS is, ISLocalToGlobalMapping ltog, ISLocalToGlobalMapping *cltog)
133: {
134:   const PetscInt *idx;
135:   PetscInt        m, *idxm, bs;

140:   ISBlockGetLocalSize(is, &m);
141:   ISBlockGetIndices(is, &idx);
142:   ISLocalToGlobalMappingGetBlockSize(ltog, &bs);
143:   PetscMalloc1(m, &idxm);
144:   if (ltog) {
145:     ISLocalToGlobalMappingApplyBlock(ltog, m, idx, idxm);
146:   } else {
147:     PetscArraycpy(idxm, idx, m);
148:   }
149:   ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is), bs, m, idxm, PETSC_OWN_POINTER, cltog);
150:   ISBlockRestoreIndices(is, &idx);
151:   return 0;
152: }

154: static PetscErrorCode MatDestroy_LocalRef(Mat B)
155: {
156:   PetscFree(B->data);
157:   return 0;
158: }

160: /*@
161:    MatCreateLocalRef - Gets a logical reference to a local submatrix, for use in assembly, that is to set values into the matrix

163:    Not Collective

165:    Input Parameters:
166: + A - Full matrix, generally parallel
167: . isrow - Local index set for the rows
168: - iscol - Local index set for the columns

170:    Output Parameter:
171: . newmat - New serial Mat

173:    Level: developer

175:    Notes:
176:    Most will use `MatGetLocalSubMatrix()` which returns a real matrix corresponding to the local
177:    block if it available, such as with matrix formats that store these blocks separately.

179:    The new matrix forwards `MatSetValuesLocal()` and `MatSetValuesBlockedLocal()` to the global system.
180:    In general, it does not define `MatMult()` or any other functions.  Local submatrices can be nested.

182: .seealso: MATSUBMATRIX`, `MatCreateSubMatrixVirtual()`, `MatSetValuesLocal()`, `MatSetValuesBlockedLocal()`, `MatGetLocalSubMatrix()`, `MatCreateSubMatrix()`
183: @*/
184: PetscErrorCode MatCreateLocalRef(Mat A, IS isrow, IS iscol, Mat *newmat)
185: {
186:   Mat_LocalRef *lr;
187:   Mat           B;
188:   PetscInt      m, n;
189:   PetscBool     islr;

196:   *newmat = NULL;

198:   MatCreate(PETSC_COMM_SELF, &B);
199:   ISGetLocalSize(isrow, &m);
200:   ISGetLocalSize(iscol, &n);
201:   MatSetSizes(B, m, n, m, n);
202:   PetscObjectChangeTypeName((PetscObject)B, MATLOCALREF);
203:   MatSetUp(B);

205:   B->ops->destroy = MatDestroy_LocalRef;

207:   PetscNew(&lr);
208:   B->data = (void *)lr;

210:   PetscObjectTypeCompare((PetscObject)A, MATLOCALREF, &islr);
211:   if (islr) {
212:     Mat_LocalRef *alr = (Mat_LocalRef *)A->data;
213:     lr->Top           = alr->Top;
214:   } else {
215:     /* This does not increase the reference count because MatLocalRef is not allowed to live longer than its parent */
216:     lr->Top = A;
217:   }
218:   {
219:     ISLocalToGlobalMapping rltog, cltog;
220:     PetscInt               arbs, acbs, rbs, cbs;

222:     /* We will translate directly to global indices for the top level */
223:     lr->SetValues        = MatSetValues;
224:     lr->SetValuesBlocked = MatSetValuesBlocked;

226:     B->ops->setvalueslocal = MatSetValuesLocal_LocalRef_Scalar;

228:     ISL2GCompose(isrow, A->rmap->mapping, &rltog);
229:     if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) {
230:       PetscObjectReference((PetscObject)rltog);
231:       cltog = rltog;
232:     } else {
233:       ISL2GCompose(iscol, A->cmap->mapping, &cltog);
234:     }
235:     /* Remember if the ISes we used to pull out the submatrix are of type ISBLOCK.  Will be used later in
236:      * MatSetValuesLocal_LocalRef_Scalar. */
237:     PetscObjectTypeCompare((PetscObject)isrow, ISBLOCK, &lr->rowisblock);
238:     PetscObjectTypeCompare((PetscObject)iscol, ISBLOCK, &lr->colisblock);
239:     MatSetLocalToGlobalMapping(B, rltog, cltog);
240:     ISLocalToGlobalMappingDestroy(&rltog);
241:     ISLocalToGlobalMappingDestroy(&cltog);

243:     MatGetBlockSizes(A, &arbs, &acbs);
244:     ISGetBlockSize(isrow, &rbs);
245:     ISGetBlockSize(iscol, &cbs);
246:     /* Always support block interface insertion on submatrix */
247:     PetscLayoutSetBlockSize(B->rmap, rbs);
248:     PetscLayoutSetBlockSize(B->cmap, cbs);
249:     if (arbs != rbs || acbs != cbs || (arbs == 1 && acbs == 1)) {
250:       /* Top-level matrix has different block size, so we have to call its scalar insertion interface */
251:       B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Scalar;
252:     } else {
253:       /* Block sizes match so we can forward values to the top level using the block interface */
254:       B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Block;

256:       ISL2GComposeBlock(isrow, A->rmap->mapping, &rltog);
257:       if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) {
258:         PetscObjectReference((PetscObject)rltog);
259:         cltog = rltog;
260:       } else {
261:         ISL2GComposeBlock(iscol, A->cmap->mapping, &cltog);
262:       }
263:       MatSetLocalToGlobalMapping(B, rltog, cltog);
264:       ISLocalToGlobalMappingDestroy(&rltog);
265:       ISLocalToGlobalMappingDestroy(&cltog);
266:     }
267:   }
268:   *newmat = B;
269:   return 0;
270: }