Actual source code: shellcnv.c

      1: #include <petsc/private/matimpl.h>
      2: #include <petsc/private/vecimpl.h>

      4: PetscErrorCode MatConvert_Shell(Mat oldmat, MatType newtype, MatReuse reuse, Mat *newmat)
      5: {
      6:   Mat          mat;
      7:   Vec          in, out;
      8:   PetscScalar *array;
      9:   PetscInt    *dnnz, *onnz, *dnnzu, *onnzu;
     10:   PetscInt     cst, cen, Nbs, mbs, nbs, rbs, cbs;
     11:   PetscInt     im, i, m, n, M, N, *rows, start;

     13:   PetscFunctionBegin;
     14:   PetscCall(MatGetOwnershipRange(oldmat, &start, NULL));
     15:   PetscCall(MatGetOwnershipRangeColumn(oldmat, &cst, &cen));
     16:   PetscCall(MatCreateVecs(oldmat, &in, &out));
     17:   PetscCall(MatGetLocalSize(oldmat, &m, &n));
     18:   PetscCall(MatGetSize(oldmat, &M, &N));
     19:   PetscCall(PetscMalloc1(m, &rows));
     20:   if (reuse != MAT_REUSE_MATRIX) {
     21:     PetscCall(MatCreate(PetscObjectComm((PetscObject)oldmat), &mat));
     22:     PetscCall(MatSetSizes(mat, m, n, M, N));
     23:     PetscCall(MatSetType(mat, newtype));
     24:     PetscCall(MatSetBlockSizesFromMats(mat, oldmat, oldmat));
     25:     PetscCall(MatGetBlockSizes(mat, &rbs, &cbs));
     26:     mbs = m / rbs;
     27:     nbs = n / cbs;
     28:     Nbs = N / cbs;
     29:     cst = cst / cbs;
     30:     PetscCall(PetscMalloc4(mbs, &dnnz, mbs, &onnz, mbs, &dnnzu, mbs, &onnzu));
     31:     for (i = 0; i < mbs; i++) {
     32:       dnnz[i]  = nbs;
     33:       onnz[i]  = Nbs - nbs;
     34:       dnnzu[i] = PetscMax(nbs - i, 0);
     35:       onnzu[i] = PetscMax(Nbs - (cst + nbs), 0);
     36:     }
     37:     PetscCall(MatXAIJSetPreallocation(mat, PETSC_DECIDE, dnnz, onnz, dnnzu, onnzu));
     38:     PetscCall(PetscFree4(dnnz, onnz, dnnzu, onnzu));
     39:     PetscCall(VecSetOption(in, VEC_IGNORE_OFF_PROC_ENTRIES, PETSC_TRUE));
     40:     PetscCall(MatSetUp(mat));
     41:   } else {
     42:     mat = *newmat;
     43:     PetscCall(MatZeroEntries(mat));
     44:   }
     45:   for (i = 0; i < N; i++) {
     46:     PetscInt j;

     48:     PetscCall(VecZeroEntries(in));
     49:     if (in->ops->setvalues) PetscCall(VecSetValue(in, i, 1., INSERT_VALUES));
     50:     else {
     51:       if (i >= cst && i < cen) {
     52:         PetscCall(VecGetArray(in, &array));
     53:         array[i - cst] = 1.0;
     54:         PetscCall(VecRestoreArray(in, &array));
     55:       }
     56:     }
     57:     PetscCall(VecAssemblyBegin(in));
     58:     PetscCall(VecAssemblyEnd(in));
     59:     PetscCall(MatMult(oldmat, in, out));
     60:     PetscCall(VecGetArray(out, &array));
     61:     for (j = 0, im = 0; j < m; j++) {
     62:       if (PetscAbsScalar(array[j]) == 0.0) continue;
     63:       rows[im]  = j + start;
     64:       array[im] = array[j];
     65:       im++;
     66:     }
     67:     PetscCall(MatSetValues(mat, im, rows, 1, &i, array, INSERT_VALUES));
     68:     PetscCall(VecRestoreArray(out, &array));
     69:   }
     70:   PetscCall(PetscFree(rows));
     71:   PetscCall(VecDestroy(&in));
     72:   PetscCall(VecDestroy(&out));
     73:   PetscCall(MatAssemblyBegin(mat, MAT_FINAL_ASSEMBLY));
     74:   PetscCall(MatAssemblyEnd(mat, MAT_FINAL_ASSEMBLY));
     75:   if (reuse == MAT_INPLACE_MATRIX) {
     76:     PetscCall(MatHeaderReplace(oldmat, &mat));
     77:   } else {
     78:     *newmat = mat;
     79:   }
     80:   PetscFunctionReturn(PETSC_SUCCESS);
     81: }

     83: static PetscErrorCode MatGetDiagonal_CF(Mat A, Vec X)
     84: {
     85:   Mat B;

     87:   PetscFunctionBegin;
     88:   PetscCall(MatShellGetContext(A, &B));
     89:   PetscCheck(B, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing user matrix");
     90:   PetscCall(MatGetDiagonal(B, X));
     91:   PetscFunctionReturn(PETSC_SUCCESS);
     92: }

     94: static PetscErrorCode MatMult_CF(Mat A, Vec X, Vec Y)
     95: {
     96:   Mat B;

     98:   PetscFunctionBegin;
     99:   PetscCall(MatShellGetContext(A, &B));
    100:   PetscCheck(B, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing user matrix");
    101:   PetscCall(MatMult(B, X, Y));
    102:   PetscFunctionReturn(PETSC_SUCCESS);
    103: }

    105: static PetscErrorCode MatMultTranspose_CF(Mat A, Vec X, Vec Y)
    106: {
    107:   Mat B;

    109:   PetscFunctionBegin;
    110:   PetscCall(MatShellGetContext(A, &B));
    111:   PetscCheck(B, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing user matrix");
    112:   PetscCall(MatMultTranspose(B, X, Y));
    113:   PetscFunctionReturn(PETSC_SUCCESS);
    114: }

    116: static PetscErrorCode MatDestroy_CF(Mat A)
    117: {
    118:   Mat B;

    120:   PetscFunctionBegin;
    121:   PetscCall(MatShellGetContext(A, &B));
    122:   PetscCheck(B, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing user matrix");
    123:   PetscCall(MatDestroy(&B));
    124:   PetscCall(MatShellSetContext(A, NULL));
    125:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_anytype_C", NULL));
    126:   PetscFunctionReturn(PETSC_SUCCESS);
    127: }

    129: typedef struct {
    130:   void              *ctx;
    131:   PetscCtxDestroyFn *destroy;
    132:   PetscErrorCode (*numeric)(Mat);
    133:   MatProductType ptype;
    134:   Mat            Dwork;
    135: } MatProductCtx_MatMatCF;

    137: static PetscErrorCode MatProductDestroy_CF(PetscCtxRt data)
    138: {
    139:   MatProductCtx_MatMatCF *mmcfdata = *(MatProductCtx_MatMatCF **)data;

    141:   PetscFunctionBegin;
    142:   if (mmcfdata->destroy) PetscCall((*mmcfdata->destroy)(&mmcfdata->ctx));
    143:   PetscCall(MatDestroy(&mmcfdata->Dwork));
    144:   PetscCall(PetscFree(mmcfdata));
    145:   PetscFunctionReturn(PETSC_SUCCESS);
    146: }

    148: static PetscErrorCode MatProductNumericPhase_CF(Mat A, Mat B, Mat C, void *data)
    149: {
    150:   MatProductCtx_MatMatCF *mmcfdata = (MatProductCtx_MatMatCF *)data;

    152:   PetscFunctionBegin;
    153:   PetscCheck(mmcfdata, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Missing data");
    154:   PetscCheck(mmcfdata->numeric, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Missing numeric operation");
    155:   /* the MATSHELL interface allows us to play with the product data */
    156:   PetscCall(PetscNew(&C->product));
    157:   C->product->type  = mmcfdata->ptype;
    158:   C->product->data  = mmcfdata->ctx;
    159:   C->product->Dwork = mmcfdata->Dwork;
    160:   PetscCall(MatShellGetContext(A, &C->product->A));
    161:   C->product->B = B;
    162:   PetscCall((*mmcfdata->numeric)(C));
    163:   PetscCall(PetscFree(C->product));
    164:   PetscFunctionReturn(PETSC_SUCCESS);
    165: }

    167: static PetscErrorCode MatProductSymbolicPhase_CF(Mat A, Mat B, Mat C, void **data)
    168: {
    169:   MatProductCtx_MatMatCF *mmcfdata;

    171:   PetscFunctionBegin;
    172:   PetscCall(MatShellGetContext(A, &C->product->A));
    173:   PetscCall(MatProductSetFromOptions(C));
    174:   PetscCall(MatProductSymbolic(C));
    175:   /* the MATSHELL interface does not allow non-empty product data */
    176:   PetscCall(PetscNew(&mmcfdata));

    178:   mmcfdata->numeric = C->ops->productnumeric;
    179:   mmcfdata->ptype   = C->product->type;
    180:   mmcfdata->ctx     = C->product->data;
    181:   mmcfdata->destroy = C->product->destroy;
    182:   mmcfdata->Dwork   = C->product->Dwork;

    184:   C->product->Dwork   = NULL;
    185:   C->product->data    = NULL;
    186:   C->product->destroy = NULL;
    187:   C->product->A       = A;

    189:   *data = mmcfdata;
    190:   PetscFunctionReturn(PETSC_SUCCESS);
    191: }

    193: /* only for A of type shell, mainly used for MatMat operations of shells with AXPYs */
    194: static PetscErrorCode MatProductSetFromOptions_CF(Mat D)
    195: {
    196:   Mat A, B, Ain;
    197:   PetscErrorCode (*Af)(Mat) = NULL;
    198:   PetscBool flg;

    200:   PetscFunctionBegin;
    201:   MatCheckProduct(D, 1);
    202:   if (D->product->type == MATPRODUCT_ABC) PetscFunctionReturn(PETSC_SUCCESS);
    203:   A = D->product->A;
    204:   B = D->product->B;
    205:   PetscCall(MatIsShell(A, &flg));
    206:   if (!flg) PetscFunctionReturn(PETSC_SUCCESS);
    207:   PetscCall(PetscObjectQueryFunction((PetscObject)A, "MatProductSetFromOptions_anytype_C", &Af));
    208:   if (Af == MatProductSetFromOptions_CF) PetscCall(MatShellGetContext(A, &Ain));
    209:   else PetscFunctionReturn(PETSC_SUCCESS);
    210:   D->product->A = Ain;
    211:   PetscCall(MatProductSetFromOptions(D));
    212:   D->product->A = A;
    213:   if (D->ops->productsymbolic) { /* we have a symbolic match, now populate the MATSHELL operations */
    214:     PetscCall(MatShellSetMatProductOperation(A, D->product->type, MatProductSymbolicPhase_CF, MatProductNumericPhase_CF, MatProductDestroy_CF, ((PetscObject)B)->type_name, NULL));
    215:     PetscCall(MatProductSetFromOptions(D));
    216:   }
    217:   PetscFunctionReturn(PETSC_SUCCESS);
    218: }

    220: PetscErrorCode MatConvertFrom_Shell(Mat A, MatType newtype, MatReuse reuse, Mat *B)
    221: {
    222:   Mat       M;
    223:   PetscBool flg;

    225:   PetscFunctionBegin;
    226:   PetscCall(PetscStrcmp(newtype, MATSHELL, &flg));
    227:   PetscCheck(flg, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only conversion to MATSHELL");
    228:   if (reuse == MAT_INITIAL_MATRIX) {
    229:     PetscCall(PetscObjectReference((PetscObject)A));
    230:     PetscCall(MatCreateShell(PetscObjectComm((PetscObject)A), A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N, A, &M));
    231:     PetscCall(MatSetBlockSizesFromMats(M, A, A));
    232:     PetscCall(MatShellSetOperation(M, MATOP_MULT, (PetscErrorCodeFn *)MatMult_CF));
    233:     PetscCall(MatShellSetOperation(M, MATOP_MULT_TRANSPOSE, (PetscErrorCodeFn *)MatMultTranspose_CF));
    234:     PetscCall(MatShellSetOperation(M, MATOP_GET_DIAGONAL, (PetscErrorCodeFn *)MatGetDiagonal_CF));
    235:     PetscCall(MatShellSetOperation(M, MATOP_DESTROY, (PetscErrorCodeFn *)MatDestroy_CF));
    236:     PetscCall(PetscObjectComposeFunction((PetscObject)M, "MatProductSetFromOptions_anytype_C", MatProductSetFromOptions_CF));
    237:     PetscCall(PetscFree(M->defaultvectype));
    238:     PetscCall(PetscStrallocpy(A->defaultvectype, &M->defaultvectype));
    239: #if defined(PETSC_HAVE_DEVICE)
    240:     PetscCall(MatBindToCPU(M, A->boundtocpu));
    241: #endif
    242:     *B = M;
    243:   } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "Not implemented");
    244:   PetscFunctionReturn(PETSC_SUCCESS);
    245: }