1 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 */ 7 8 #include <../src/mat/impls/aij/seq/aij.h> /*I "petscmat.h" I*/ 9 #include <../src/mat/impls/dense/seq/dense.h> 10 11 PetscErrorCode MatDestroy_SeqDense_MatTransMatMult(void *data) { 12 Mat_MatTransMatMult *atb = (Mat_MatTransMatMult *)data; 13 14 PetscFunctionBegin; 15 PetscCall(MatDestroy(&atb->mA)); 16 PetscCall(VecDestroy(&atb->bt)); 17 PetscCall(VecDestroy(&atb->ct)); 18 PetscCall(PetscFree(atb)); 19 PetscFunctionReturn(0); 20 } 21 22 static PetscErrorCode MatTMatTMultNumeric_SeqAIJ_SeqDense(Mat, Mat, Mat); 23 24 PETSC_INTERN PetscErrorCode MatTMatTMultSymbolic_SeqAIJ_SeqDense(Mat A, Mat B, PetscReal fill, Mat C) { 25 Mat_MatTransMatMult *atb; 26 PetscBool cisdense; 27 PetscInt dofm; 28 29 PetscFunctionBegin; 30 MatCheckProduct(C, 4); 31 PetscCheck(!C->product->data, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Extra product struct not empty"); 32 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]); 33 34 /* create output dense matrix C */ 35 if (C->product->type == MATPRODUCT_AtB) { 36 PetscCall(MatSetSizes(C, A->cmap->n, B->cmap->N, A->cmap->n, B->cmap->N)); 37 dofm = B->cmap->n; 38 } else { 39 PetscCall(MatSetSizes(C, A->rmap->n, B->rmap->N, A->rmap->n, B->rmap->N)); 40 dofm = B->rmap->n; 41 } 42 PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATSEQDENSE, MATSEQDENSECUDA, "")); 43 if (!cisdense) PetscCall(MatSetType(C, ((PetscObject)B)->type_name)); 44 PetscCall(MatSetUp(C)); 45 46 /* create additional data structure for the product */ 47 PetscCall(PetscNew(&atb)); 48 PetscCall(MatCreateMAIJ(A, dofm, &atb->mA)); 49 PetscCall(MatCreateVecs(atb->mA, &atb->ct, &atb->bt)); 50 C->product->data = atb; 51 C->product->destroy = MatDestroy_SeqDense_MatTransMatMult; 52 53 if (C->product->type == MATPRODUCT_AtB) { 54 C->ops->transposematmultnumeric = MatTMatTMultNumeric_SeqAIJ_SeqDense; 55 } else { 56 C->ops->mattransposemultnumeric = MatTMatTMultNumeric_SeqAIJ_SeqDense; 57 } 58 PetscFunctionReturn(0); 59 } 60 61 PetscErrorCode MatTMatTMultNumeric_SeqAIJ_SeqDense(Mat A, Mat B, Mat C) { 62 PetscInt i, j, m = A->rmap->n, n = A->cmap->n, blda, clda; 63 PetscInt mdof = C->cmap->N; 64 const PetscScalar *Barray; 65 PetscScalar *Carray; 66 Mat_MatTransMatMult *atb; 67 Vec bt, ct; 68 69 PetscFunctionBegin; 70 MatCheckProduct(C, 3); 71 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]); 72 atb = (Mat_MatTransMatMult *)C->product->data; 73 PetscCheck(atb, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing product struct"); 74 bt = atb->bt; 75 ct = atb->ct; 76 77 PetscCall(MatDenseGetArrayRead(B, &Barray)); 78 PetscCall(MatDenseGetLDA(B, &blda)); 79 PetscCall(MatDenseGetArrayWrite(C, &Carray)); 80 PetscCall(MatDenseGetLDA(C, &clda)); 81 if (C->product->type == MATPRODUCT_AtB) { /* transpose local array of B, then copy it to vector bt */ 82 const PetscScalar *ctarray; 83 PetscScalar *btarray; 84 85 PetscCall(VecGetArrayWrite(bt, &btarray)); 86 for (j = 0; j < mdof; j++) { 87 for (i = 0; i < m; i++) btarray[i * mdof + j] = Barray[j * blda + i]; 88 } 89 PetscCall(VecRestoreArrayWrite(bt, &btarray)); 90 91 /* compute ct = mA^T * cb */ 92 PetscCall(MatMultTranspose(atb->mA, bt, ct)); 93 94 /* transpose local array of ct to matrix C */ 95 PetscCall(VecGetArrayRead(ct, &ctarray)); 96 for (j = 0; j < mdof; j++) { 97 for (i = 0; i < n; i++) Carray[j * clda + i] = ctarray[i * mdof + j]; 98 } 99 PetscCall(VecRestoreArrayRead(ct, &ctarray)); 100 } else { 101 const PetscScalar *btarray; 102 PetscScalar *ctarray; 103 104 if (blda == B->rmap->n) { 105 PetscCall(VecPlaceArray(ct, Barray)); 106 } else { 107 PetscInt bn = B->cmap->n; 108 PetscInt bm = B->rmap->n; 109 110 PetscCall(VecGetArrayWrite(ct, &ctarray)); 111 for (j = 0; j < bn; j++) { 112 for (i = 0; i < bm; i++) ctarray[j * bm + i] = Barray[j * blda + i]; 113 } 114 PetscCall(VecRestoreArrayWrite(ct, &ctarray)); 115 } 116 117 PetscCall(MatMult(atb->mA, ct, bt)); 118 if (blda == B->rmap->n) PetscCall(VecResetArray(ct)); 119 PetscCall(VecGetArrayRead(bt, &btarray)); 120 for (j = 0; j < mdof; j++) { 121 for (i = 0; i < m; i++) Carray[j * clda + i] = btarray[i * mdof + j]; 122 } 123 PetscCall(VecRestoreArrayRead(bt, &btarray)); 124 } 125 PetscCall(MatDenseRestoreArrayRead(B, &Barray)); 126 PetscCall(MatDenseRestoreArray(C, &Carray)); 127 PetscFunctionReturn(0); 128 } 129