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) PetscCall(ISLocalToGlobalMappingApply(ltog, m, idx, idxm));
    126:   else PetscCall(PetscArraycpy(idxm, idx, m));
    127:   PetscCall(ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is), bs, m, idxm, PETSC_OWN_POINTER, cltog));
    128:   PetscCall(ISRestoreIndices(is, &idx));
    129:   PetscFunctionReturn(PETSC_SUCCESS);
    130: }

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

    137:   PetscFunctionBegin;
    140:   PetscAssertPointer(cltog, 3);
    141:   PetscCall(ISBlockGetLocalSize(is, &m));
    142:   PetscCall(ISBlockGetIndices(is, &idx));
    143:   PetscCall(ISLocalToGlobalMappingGetBlockSize(ltog, &bs));
    144:   PetscCall(PetscMalloc1(m, &idxm));
    145:   if (ltog) PetscCall(ISLocalToGlobalMappingApplyBlock(ltog, m, idx, idxm));
    146:   else PetscCall(PetscArraycpy(idxm, idx, m));
    147:   PetscCall(ISLocalToGlobalMappingCreate(PetscObjectComm((PetscObject)is), bs, m, idxm, PETSC_OWN_POINTER, cltog));
    148:   PetscCall(ISBlockRestoreIndices(is, &idx));
    149:   PetscFunctionReturn(PETSC_SUCCESS);
    150: }

    152: static PetscErrorCode MatZeroRowsLocal_LocalRef(Mat A, PetscInt n, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
    153: {
    154:   PetscInt     *rows_l;
    155:   Mat_LocalRef *lr = (Mat_LocalRef *)A->data;

    157:   PetscFunctionBegin;
    158:   PetscCall(PetscMalloc1(n, &rows_l));
    159:   PetscCall(ISLocalToGlobalMappingApply(A->rmap->mapping, n, rows, rows_l));
    160:   PetscCall(MatZeroRows(lr->Top, n, rows_l, diag, x, b));
    161:   PetscCall(PetscFree(rows_l));
    162:   PetscFunctionReturn(PETSC_SUCCESS);
    163: }

    165: static PetscErrorCode MatZeroRowsColumnsLocal_LocalRef(Mat A, PetscInt n, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
    166: {
    167:   PetscInt     *rows_l;
    168:   Mat_LocalRef *lr = (Mat_LocalRef *)A->data;

    170:   PetscFunctionBegin;
    171:   PetscCall(PetscMalloc1(n, &rows_l));
    172:   PetscCall(ISLocalToGlobalMappingApply(A->rmap->mapping, n, rows, rows_l));
    173:   PetscCall(MatZeroRowsColumns(lr->Top, n, rows_l, diag, x, b));
    174:   PetscCall(PetscFree(rows_l));
    175:   PetscFunctionReturn(PETSC_SUCCESS);
    176: }

    178: static PetscErrorCode MatDestroy_LocalRef(Mat B)
    179: {
    180:   PetscFunctionBegin;
    181:   PetscCall(PetscFree(B->data));
    182:   PetscFunctionReturn(PETSC_SUCCESS);
    183: }

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

    188:   Not Collective

    190:   Input Parameters:
    191: + A     - full matrix, generally parallel
    192: . isrow - Local index set for the rows
    193: - iscol - Local index set for the columns

    195:   Output Parameter:
    196: . newmat - new serial `Mat`

    198:   Level: developer

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

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

    207: .seealso: [](ch_matrices), `Mat`, `MATSUBMATRIX`, `MatCreateSubMatrixVirtual()`, `MatSetValuesLocal()`, `MatSetValuesBlockedLocal()`, `MatGetLocalSubMatrix()`, `MatCreateSubMatrix()`
    208: @*/
    209: PetscErrorCode MatCreateLocalRef(Mat A, IS isrow, IS iscol, Mat *newmat)
    210: {
    211:   Mat_LocalRef *lr;
    212:   Mat           B;
    213:   PetscInt      m, n;
    214:   PetscBool     islr;

    216:   PetscFunctionBegin;
    220:   PetscAssertPointer(newmat, 4);
    221:   PetscCheck(A->rmap->mapping, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Matrix must have local to global mapping provided before this call");
    222:   *newmat = NULL;

    224:   PetscCall(MatCreate(PETSC_COMM_SELF, &B));
    225:   PetscCall(ISGetLocalSize(isrow, &m));
    226:   PetscCall(ISGetLocalSize(iscol, &n));
    227:   PetscCall(MatSetSizes(B, m, n, m, n));
    228:   PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATLOCALREF));
    229:   PetscCall(MatSetUp(B));

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

    233:   PetscCall(PetscNew(&lr));
    234:   B->data = (void *)lr;

    236:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATLOCALREF, &islr));
    237:   if (islr) {
    238:     Mat_LocalRef *alr = (Mat_LocalRef *)A->data;
    239:     lr->Top           = alr->Top;
    240:   } else {
    241:     /* This does not increase the reference count because MatLocalRef is not allowed to live longer than its parent */
    242:     lr->Top = A;
    243:   }
    244:   {
    245:     ISLocalToGlobalMapping rltog, cltog;
    246:     PetscInt               arbs, acbs, rbs, cbs;

    248:     /* We will translate directly to global indices for the top level */
    249:     lr->SetValues        = MatSetValues;
    250:     lr->SetValuesBlocked = MatSetValuesBlocked;

    252:     B->ops->setvalueslocal       = MatSetValuesLocal_LocalRef_Scalar;
    253:     B->ops->zerorowslocal        = MatZeroRowsLocal_LocalRef;
    254:     B->ops->zerorowscolumnslocal = MatZeroRowsColumnsLocal_LocalRef;

    256:     PetscCall(ISL2GCompose(isrow, A->rmap->mapping, &rltog));
    257:     if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) {
    258:       PetscCall(PetscObjectReference((PetscObject)rltog));
    259:       cltog = rltog;
    260:     } else {
    261:       PetscCall(ISL2GCompose(iscol, A->cmap->mapping, &cltog));
    262:     }
    263:     /* Remember if the ISes we used to pull out the submatrix are of type ISBLOCK.  Will be used later in
    264:      * MatSetValuesLocal_LocalRef_Scalar. */
    265:     PetscCall(PetscObjectTypeCompare((PetscObject)isrow, ISBLOCK, &lr->rowisblock));
    266:     PetscCall(PetscObjectTypeCompare((PetscObject)iscol, ISBLOCK, &lr->colisblock));
    267:     PetscCall(MatSetLocalToGlobalMapping(B, rltog, cltog));
    268:     PetscCall(ISLocalToGlobalMappingDestroy(&rltog));
    269:     PetscCall(ISLocalToGlobalMappingDestroy(&cltog));

    271:     PetscCall(MatGetBlockSizes(A, &arbs, &acbs));
    272:     PetscCall(ISGetBlockSize(isrow, &rbs));
    273:     PetscCall(ISGetBlockSize(iscol, &cbs));
    274:     /* Always support block interface insertion on submatrix */
    275:     PetscCall(PetscLayoutSetBlockSize(B->rmap, rbs));
    276:     PetscCall(PetscLayoutSetBlockSize(B->cmap, cbs));
    277:     if (arbs != rbs || acbs != cbs || (arbs == 1 && acbs == 1)) {
    278:       /* Top-level matrix has different block size, so we have to call its scalar insertion interface */
    279:       B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Scalar;
    280:     } else {
    281:       /* Block sizes match so we can forward values to the top level using the block interface */
    282:       B->ops->setvaluesblockedlocal = MatSetValuesBlockedLocal_LocalRef_Block;

    284:       PetscCall(ISL2GComposeBlock(isrow, A->rmap->mapping, &rltog));
    285:       if (isrow == iscol && A->rmap->mapping == A->cmap->mapping) {
    286:         PetscCall(PetscObjectReference((PetscObject)rltog));
    287:         cltog = rltog;
    288:       } else {
    289:         PetscCall(ISL2GComposeBlock(iscol, A->cmap->mapping, &cltog));
    290:       }
    291:       PetscCall(MatSetLocalToGlobalMapping(B, rltog, cltog));
    292:       PetscCall(ISLocalToGlobalMappingDestroy(&rltog));
    293:       PetscCall(ISLocalToGlobalMappingDestroy(&cltog));
    294:     }
    295:   }
    296:   *newmat = B;
    297:   PetscFunctionReturn(PETSC_SUCCESS);
    298: }