1: /*
  2:   Defines matrix-matrix product routines for
  3:           C = A^T * B and C = A * B^t
  4:   with A SeqAIJ and B SeqDense
  5: */

  7: #include <../src/mat/impls/aij/seq/aij.h>
  8: #include <../src/mat/impls/dense/seq/dense.h>

 10: static PetscErrorCode MatProductCtxDestroy_SeqDense_MatTransMatMult(PetscCtxRt data)
 11: {
 12:   MatProductCtx_MatTransMatMult *atb = *(MatProductCtx_MatTransMatMult **)data;

 14:   PetscFunctionBegin;
 15:   PetscCall(MatDestroy(&atb->mA));
 16:   PetscCall(VecDestroy(&atb->bt));
 17:   PetscCall(VecDestroy(&atb->ct));
 18:   PetscCall(PetscFree(atb));
 19:   PetscFunctionReturn(PETSC_SUCCESS);
 20: }

 22: static PetscErrorCode MatTMatTMultNumeric_SeqAIJ_SeqDense(Mat, Mat, Mat);

 24: PETSC_INTERN PetscErrorCode MatTMatTMultSymbolic_SeqAIJ_SeqDense(Mat A, Mat B, PetscReal fill, Mat C)
 25: {
 26:   MatProductCtx_MatTransMatMult *atb;
 27:   PetscBool                      cisdense;
 28:   PetscInt                       dofm;

 30:   PetscFunctionBegin;
 31:   MatCheckProduct(C, 4);
 32:   PetscCheck(!C->product->data, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Extra product struct not empty");
 33:   PetscCheck(C->product->type == MATPRODUCT_ABt || C->product->type == MATPRODUCT_AtB, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Not for product type %s", MatProductTypes[C->product->type]);

 35:   /* create output dense matrix C */
 36:   if (C->product->type == MATPRODUCT_AtB) {
 37:     PetscCall(MatSetSizes(C, A->cmap->n, B->cmap->N, A->cmap->n, B->cmap->N));
 38:     dofm = B->cmap->n;
 39:   } else {
 40:     PetscCall(MatSetSizes(C, A->rmap->n, B->rmap->N, A->rmap->n, B->rmap->N));
 41:     dofm = B->rmap->n;
 42:   }
 43:   PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATSEQDENSE, MATSEQDENSECUDA, ""));
 44:   if (!cisdense) {
 45:     PetscCall(MatSetType(C, ((PetscObject)B)->type_name));
 46:     PetscCall(MatSetVecType(C, B->defaultvectype));
 47:   }
 48:   PetscCall(MatSetUp(C));

 50:   /* create additional data structure for the product */
 51:   PetscCall(PetscNew(&atb));
 52:   PetscCall(MatCreateMAIJ(A, dofm, &atb->mA));
 53:   PetscCall(MatCreateVecs(atb->mA, &atb->ct, &atb->bt));
 54:   C->product->data    = atb;
 55:   C->product->destroy = MatProductCtxDestroy_SeqDense_MatTransMatMult;

 57:   if (C->product->type == MATPRODUCT_AtB) {
 58:     C->ops->transposematmultnumeric = MatTMatTMultNumeric_SeqAIJ_SeqDense;
 59:   } else {
 60:     C->ops->mattransposemultnumeric = MatTMatTMultNumeric_SeqAIJ_SeqDense;
 61:   }
 62:   PetscFunctionReturn(PETSC_SUCCESS);
 63: }

 65: static PetscErrorCode MatTMatTMultNumeric_SeqAIJ_SeqDense(Mat A, Mat B, Mat C)
 66: {
 67:   PetscInt                       i, j, m = A->rmap->n, n = A->cmap->n, blda, clda;
 68:   PetscInt                       mdof = C->cmap->N;
 69:   const PetscScalar             *Barray;
 70:   PetscScalar                   *Carray;
 71:   MatProductCtx_MatTransMatMult *atb;
 72:   Vec                            bt, ct;

 74:   PetscFunctionBegin;
 75:   MatCheckProduct(C, 3);
 76:   PetscCheck(C->product->type == MATPRODUCT_ABt || C->product->type == MATPRODUCT_AtB, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Not for product type %s", MatProductTypes[C->product->type]);
 77:   atb = (MatProductCtx_MatTransMatMult *)C->product->data;
 78:   PetscCheck(atb, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing product struct");
 79:   bt = atb->bt;
 80:   ct = atb->ct;

 82:   PetscCall(MatDenseGetArrayRead(B, &Barray));
 83:   PetscCall(MatDenseGetLDA(B, &blda));
 84:   PetscCall(MatDenseGetArrayWrite(C, &Carray));
 85:   PetscCall(MatDenseGetLDA(C, &clda));
 86:   if (C->product->type == MATPRODUCT_AtB) { /* transpose local array of B, then copy it to vector bt */
 87:     const PetscScalar *ctarray;
 88:     PetscScalar       *btarray;

 90:     PetscCall(VecGetArrayWrite(bt, &btarray));
 91:     for (j = 0; j < mdof; j++) {
 92:       for (i = 0; i < m; i++) btarray[i * mdof + j] = Barray[j * blda + i];
 93:     }
 94:     PetscCall(VecRestoreArrayWrite(bt, &btarray));

 96:     /* compute ct = mA^T * cb */
 97:     PetscCall(MatMultTranspose(atb->mA, bt, ct));

 99:     /* transpose local array of ct to matrix C */
100:     PetscCall(VecGetArrayRead(ct, &ctarray));
101:     for (j = 0; j < mdof; j++) {
102:       for (i = 0; i < n; i++) Carray[j * clda + i] = ctarray[i * mdof + j];
103:     }
104:     PetscCall(VecRestoreArrayRead(ct, &ctarray));
105:   } else {
106:     const PetscScalar *btarray;
107:     PetscScalar       *ctarray;

109:     if (blda == B->rmap->n) {
110:       PetscCall(VecPlaceArray(ct, Barray));
111:     } else {
112:       PetscInt bn = B->cmap->n;
113:       PetscInt bm = B->rmap->n;

115:       PetscCall(VecGetArrayWrite(ct, &ctarray));
116:       for (j = 0; j < bn; j++) {
117:         for (i = 0; i < bm; i++) ctarray[j * bm + i] = Barray[j * blda + i];
118:       }
119:       PetscCall(VecRestoreArrayWrite(ct, &ctarray));
120:     }

122:     PetscCall(MatMult(atb->mA, ct, bt));
123:     if (blda == B->rmap->n) PetscCall(VecResetArray(ct));
124:     PetscCall(VecGetArrayRead(bt, &btarray));
125:     for (j = 0; j < mdof; j++) {
126:       for (i = 0; i < m; i++) Carray[j * clda + i] = btarray[i * mdof + j];
127:     }
128:     PetscCall(VecRestoreArrayRead(bt, &btarray));
129:   }
130:   PetscCall(MatDenseRestoreArrayRead(B, &Barray));
131:   PetscCall(MatDenseRestoreArray(C, &Carray));
132:   PetscFunctionReturn(PETSC_SUCCESS);
133: }