Actual source code: mattransposematmult.c

petsc-master 2020-09-24
Report Typos and Errors

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

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

 11: PetscErrorCode MatDestroy_SeqDense_MatTransMatMult(void *data)
 12: {
 13:   PetscErrorCode      ierr;
 14:   Mat_MatTransMatMult *atb = (Mat_MatTransMatMult *)data;

 17:   MatDestroy(&atb->mA);
 18:   VecDestroy(&atb->bt);
 19:   VecDestroy(&atb->ct);
 20:   PetscFree(atb);
 21:   return(0);
 22: }

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

 26: PETSC_INTERN PetscErrorCode MatTMatTMultSymbolic_SeqAIJ_SeqDense(Mat A,Mat B,PetscReal fill,Mat C)
 27: {
 28:   PetscErrorCode      ierr;
 29:   Mat_MatTransMatMult *atb;
 30:   PetscBool           cisdense;
 31:   PetscInt            dofm;

 34:   MatCheckProduct(C,4);
 35:   if (C->product->data) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty");
 36:   if (C->product->type != MATPRODUCT_ABt && C->product->type != MATPRODUCT_AtB) SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Not for product type %s",MatProductTypes[C->product->type]);

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

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

 59:   if (C->product->type == MATPRODUCT_AtB) {
 60:     C->ops->transposematmultnumeric = MatTMatTMultNumeric_SeqAIJ_SeqDense;
 61:   } else {
 62:     C->ops->mattransposemultnumeric = MatTMatTMultNumeric_SeqAIJ_SeqDense;
 63:   }
 64:   return(0);
 65: }

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

 78:   MatCheckProduct(C,3);
 79:   if (C->product->type != MATPRODUCT_ABt && C->product->type != MATPRODUCT_AtB) SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Not for product type %s",MatProductTypes[C->product->type]);
 80:   atb = (Mat_MatTransMatMult *)C->product->data;
 81:   if (!atb) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Missing product struct");
 82:   bt = atb->bt;
 83:   ct = atb->ct;

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

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

 99:     /* compute ct = mA^T * cb */
100:     MatMultTranspose(atb->mA,bt,ct);

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

112:     if (blda == B->rmap->n) {
113:       VecPlaceArray(ct,Barray);
114:     } else {
115:       PetscInt bn = B->cmap->n;
116:       PetscInt bm = B->rmap->n;

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

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