Actual source code: mpimatmatmatmult.c

      1: /*
      2:   Defines matrix-matrix-matrix product routines for MPIAIJ matrices
      3:           D = A * B * C
      4: */
      5: #include <../src/mat/impls/aij/mpi/mpiaij.h>

      7: #if defined(PETSC_HAVE_HYPRE)
      8: PETSC_INTERN PetscErrorCode MatTransposeMatMatMultSymbolic_AIJ_AIJ_AIJ_wHYPRE(Mat, Mat, Mat, PetscReal, Mat);
      9: PETSC_INTERN PetscErrorCode MatTransposeMatMatMultNumeric_AIJ_AIJ_AIJ_wHYPRE(Mat, Mat, Mat, Mat);

     11: PETSC_INTERN PetscErrorCode MatProductNumeric_ABC_Transpose_AIJ_AIJ(Mat RAP)
     12: {
     13:   Mat_Product *product = RAP->product;
     14:   Mat          Rt, R = product->A, A = product->B, P = product->C;

     16:   PetscFunctionBegin;
     17:   PetscCall(MatTransposeGetMat(R, &Rt));
     18:   PetscCall(MatTransposeMatMatMultNumeric_AIJ_AIJ_AIJ_wHYPRE(Rt, A, P, RAP));
     19:   PetscFunctionReturn(PETSC_SUCCESS);
     20: }

     22: PETSC_INTERN PetscErrorCode MatProductSymbolic_ABC_Transpose_AIJ_AIJ(Mat RAP)
     23: {
     24:   Mat_Product *product = RAP->product;
     25:   Mat          Rt, R = product->A, A = product->B, P = product->C;
     26:   PetscBool    flg;

     28:   PetscFunctionBegin;
     29:   /* local sizes of matrices will be checked by the calling subroutines */
     30:   PetscCall(MatTransposeGetMat(R, &Rt));
     31:   PetscCall(PetscObjectTypeCompareAny((PetscObject)Rt, &flg, MATSEQAIJ, MATSEQAIJMKL, MATMPIAIJ, NULL));
     32:   PetscCheck(flg, PetscObjectComm((PetscObject)Rt), PETSC_ERR_SUP, "Not for matrix type %s", ((PetscObject)Rt)->type_name);
     33:   PetscCall(MatTransposeMatMatMultSymbolic_AIJ_AIJ_AIJ_wHYPRE(Rt, A, P, product->fill, RAP));
     34:   RAP->ops->productnumeric = MatProductNumeric_ABC_Transpose_AIJ_AIJ;
     35:   PetscFunctionReturn(PETSC_SUCCESS);
     36: }

     38: PETSC_INTERN PetscErrorCode MatProductSetFromOptions_Transpose_AIJ_AIJ(Mat C)
     39: {
     40:   Mat_Product *product = C->product;

     42:   PetscFunctionBegin;
     43:   PetscCheck(product->type == MATPRODUCT_ABC, PetscObjectComm((PetscObject)C), PETSC_ERR_SUP, "MatProduct type %s is not supported for Transpose, AIJ and AIJ matrices", MatProductTypes[product->type]);
     44:   C->ops->productsymbolic = MatProductSymbolic_ABC_Transpose_AIJ_AIJ;
     45:   PetscFunctionReturn(PETSC_SUCCESS);
     46: }
     47: #endif

     49: PetscErrorCode MatMatMatMultSymbolic_MPIAIJ_MPIAIJ_MPIAIJ(Mat A, Mat B, Mat C, PetscReal fill, Mat D)
     50: {
     51:   Mat          BC;
     52:   PetscBool    scalable;
     53:   Mat_Product *product;

     55:   PetscFunctionBegin;
     56:   MatCheckProduct(D, 5);
     57:   PetscCheck(!D->product->data, PetscObjectComm((PetscObject)D), PETSC_ERR_PLIB, "Product data not empty");
     58:   product = D->product;
     59:   PetscCall(MatProductCreate(B, C, NULL, &BC));
     60:   PetscCall(MatProductSetType(BC, MATPRODUCT_AB));
     61:   PetscCall(PetscStrcmp(product->alg, "scalable", &scalable));
     62:   if (scalable) {
     63:     PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ(B, C, fill, BC));
     64:     PetscCall(MatZeroEntries(BC)); /* initialize value entries of BC */
     65:     PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ(A, BC, fill, D));
     66:   } else {
     67:     PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(B, C, fill, BC));
     68:     PetscCall(MatZeroEntries(BC)); /* initialize value entries of BC */
     69:     PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(A, BC, fill, D));
     70:   }
     71:   PetscCall(MatDestroy(&product->Dwork));
     72:   product->Dwork = BC;

     74:   D->ops->matmatmultnumeric = MatMatMatMultNumeric_MPIAIJ_MPIAIJ_MPIAIJ;
     75:   PetscFunctionReturn(PETSC_SUCCESS);
     76: }

     78: PetscErrorCode MatMatMatMultNumeric_MPIAIJ_MPIAIJ_MPIAIJ(Mat A, Mat B, Mat C, Mat D)
     79: {
     80:   Mat_Product *product;
     81:   Mat          BC;

     83:   PetscFunctionBegin;
     84:   MatCheckProduct(D, 4);
     85:   PetscCheck(D->product->data, PetscObjectComm((PetscObject)D), PETSC_ERR_PLIB, "Product data empty");
     86:   product = D->product;
     87:   BC      = product->Dwork;
     88:   PetscCall((*BC->ops->matmultnumeric)(B, C, BC));
     89:   PetscCall((*D->ops->matmultnumeric)(A, BC, D));
     90:   PetscFunctionReturn(PETSC_SUCCESS);
     91: }

     93: static PetscErrorCode MatProductCtxDestroy_MPIAIJ_RARt(PetscCtxRt data)
     94: {
     95:   MatProductCtx_RARt *rart = *(MatProductCtx_RARt **)data;

     97:   PetscFunctionBegin;
     98:   PetscCall(MatDestroy(&rart->Rt));
     99:   if (rart->destroy) PetscCall((*rart->destroy)(&rart->data));
    100:   PetscCall(PetscFree(rart));
    101:   PetscFunctionReturn(PETSC_SUCCESS);
    102: }

    104: PetscErrorCode MatProductNumeric_RARt_MPIAIJ_MPIAIJ(Mat C)
    105: {
    106:   MatProductCtx_RARt *rart;
    107:   Mat                 A, R, Rt;

    109:   PetscFunctionBegin;
    110:   MatCheckProduct(C, 1);
    111:   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
    112:   rart = (MatProductCtx_RARt *)C->product->data;
    113:   A    = C->product->A;
    114:   R    = C->product->B;
    115:   Rt   = rart->Rt;
    116:   PetscCall(MatTranspose(R, MAT_REUSE_MATRIX, &Rt));
    117:   if (rart->data) C->product->data = rart->data;
    118:   PetscCall((*C->ops->matmatmultnumeric)(R, A, Rt, C));
    119:   C->product->data = rart;
    120:   PetscFunctionReturn(PETSC_SUCCESS);
    121: }

    123: PetscErrorCode MatProductSymbolic_RARt_MPIAIJ_MPIAIJ(Mat C)
    124: {
    125:   Mat                 A, R, Rt;
    126:   MatProductCtx_RARt *rart;

    128:   PetscFunctionBegin;
    129:   MatCheckProduct(C, 1);
    130:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
    131:   A = C->product->A;
    132:   R = C->product->B;
    133:   PetscCall(MatTranspose(R, MAT_INITIAL_MATRIX, &Rt));
    134:   /* product->Dwork is used to store A*Rt in MatMatMatMultSymbolic_MPIAIJ_MPIAIJ_MPIAIJ() */
    135:   PetscCall(MatMatMatMultSymbolic_MPIAIJ_MPIAIJ_MPIAIJ(R, A, Rt, C->product->fill, C));
    136:   C->ops->productnumeric = MatProductNumeric_RARt_MPIAIJ_MPIAIJ;

    138:   /* create a supporting struct */
    139:   PetscCall(PetscNew(&rart));
    140:   rart->Rt            = Rt;
    141:   rart->data          = C->product->data;
    142:   rart->destroy       = C->product->destroy;
    143:   C->product->data    = rart;
    144:   C->product->destroy = MatProductCtxDestroy_MPIAIJ_RARt;
    145:   PetscFunctionReturn(PETSC_SUCCESS);
    146: }