1ed3cc1f0SBarry Smith /* 2ed3cc1f0SBarry Smith Basic functions for basic parallel dense matrices. 347d993e7Ssuyashtn Portions of this code are under: 447d993e7Ssuyashtn Copyright (c) 2022 Advanced Micro Devices, Inc. All rights reserved. 5ed3cc1f0SBarry Smith */ 6ed3cc1f0SBarry Smith 7c6db04a5SJed Brown #include <../src/mat/impls/dense/mpi/mpidense.h> /*I "petscmat.h" I*/ 88949adfdSHong Zhang #include <../src/mat/impls/aij/mpi/mpiaij.h> 9baa3c1c6SHong Zhang #include <petscblaslapack.h> 104186a4bbSPierre Jolivet #include <petsc/private/vecimpl.h> 118965ea79SLois Curfman McInnes 12ab92ecdeSBarry Smith /*@ 1311a5261eSBarry Smith MatDenseGetLocalMatrix - For a `MATMPIDENSE` or `MATSEQDENSE` matrix returns the sequential 14ab92ecdeSBarry Smith matrix that represents the operator. For sequential matrices it returns itself. 15ab92ecdeSBarry Smith 16ab92ecdeSBarry Smith Input Parameter: 172ef1f0ffSBarry Smith . A - the sequential or MPI `MATDENSE` matrix 18ab92ecdeSBarry Smith 19ab92ecdeSBarry Smith Output Parameter: 20ab92ecdeSBarry Smith . B - the inner matrix 21ab92ecdeSBarry Smith 228e6c10adSSatish Balay Level: intermediate 238e6c10adSSatish Balay 241cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MATMPIDENSE`, `MATSEQDENSE` 25ab92ecdeSBarry Smith @*/ 26d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetLocalMatrix(Mat A, Mat *B) 27d71ae5a4SJacob Faibussowitsch { 28ab92ecdeSBarry Smith Mat_MPIDense *mat = (Mat_MPIDense *)A->data; 29ace3abfcSBarry Smith PetscBool flg; 30ab92ecdeSBarry Smith 31ab92ecdeSBarry Smith PetscFunctionBegin; 32d5ea218eSStefano Zampini PetscValidHeaderSpecific(A, MAT_CLASSID, 1); 334f572ea9SToby Isaac PetscAssertPointer(B, 2); 349566063dSJacob Faibussowitsch PetscCall(PetscObjectBaseTypeCompare((PetscObject)A, MATMPIDENSE, &flg)); 352205254eSKarl Rupp if (flg) *B = mat->A; 36f140bc17SStefano Zampini else { 379566063dSJacob Faibussowitsch PetscCall(PetscObjectBaseTypeCompare((PetscObject)A, MATSEQDENSE, &flg)); 3828b400f6SJacob Faibussowitsch PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "Not for matrix type %s", ((PetscObject)A)->type_name); 39f140bc17SStefano Zampini *B = A; 40f140bc17SStefano Zampini } 413ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 42ab92ecdeSBarry Smith } 43ab92ecdeSBarry Smith 4466976f2fSJacob Faibussowitsch static PetscErrorCode MatCopy_MPIDense(Mat A, Mat B, MatStructure s) 45d71ae5a4SJacob Faibussowitsch { 46a76f77c3SStefano Zampini Mat_MPIDense *Amat = (Mat_MPIDense *)A->data; 47a76f77c3SStefano Zampini Mat_MPIDense *Bmat = (Mat_MPIDense *)B->data; 48a76f77c3SStefano Zampini 49a76f77c3SStefano Zampini PetscFunctionBegin; 509566063dSJacob Faibussowitsch PetscCall(MatCopy(Amat->A, Bmat->A, s)); 513ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 52a76f77c3SStefano Zampini } 53a76f77c3SStefano Zampini 54d71ae5a4SJacob Faibussowitsch PetscErrorCode MatShift_MPIDense(Mat A, PetscScalar alpha) 55d71ae5a4SJacob Faibussowitsch { 562f605a99SJose E. Roman Mat_MPIDense *mat = (Mat_MPIDense *)A->data; 572f605a99SJose E. Roman PetscInt j, lda, rstart = A->rmap->rstart, rend = A->rmap->rend, rend2; 582f605a99SJose E. Roman PetscScalar *v; 592f605a99SJose E. Roman 602f605a99SJose E. Roman PetscFunctionBegin; 612f605a99SJose E. Roman PetscCall(MatDenseGetArray(mat->A, &v)); 622f605a99SJose E. Roman PetscCall(MatDenseGetLDA(mat->A, &lda)); 632f605a99SJose E. Roman rend2 = PetscMin(rend, A->cmap->N); 642f605a99SJose E. Roman if (rend2 > rstart) { 652f605a99SJose E. Roman for (j = rstart; j < rend2; j++) v[j - rstart + j * lda] += alpha; 662f605a99SJose E. Roman PetscCall(PetscLogFlops(rend2 - rstart)); 672f605a99SJose E. Roman } 682f605a99SJose E. Roman PetscCall(MatDenseRestoreArray(mat->A, &v)); 693ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 702f605a99SJose E. Roman } 712f605a99SJose E. Roman 7266976f2fSJacob Faibussowitsch static PetscErrorCode MatGetRow_MPIDense(Mat A, PetscInt row, PetscInt *nz, PetscInt **idx, PetscScalar **v) 73d71ae5a4SJacob Faibussowitsch { 74ba8c8a56SBarry Smith Mat_MPIDense *mat = (Mat_MPIDense *)A->data; 75d0f46423SBarry Smith PetscInt lrow, rstart = A->rmap->rstart, rend = A->rmap->rend; 76ba8c8a56SBarry Smith 77ba8c8a56SBarry Smith PetscFunctionBegin; 78aed4548fSBarry Smith PetscCheck(row >= rstart && row < rend, PETSC_COMM_SELF, PETSC_ERR_SUP, "only local rows"); 79ba8c8a56SBarry Smith lrow = row - rstart; 809566063dSJacob Faibussowitsch PetscCall(MatGetRow(mat->A, lrow, nz, (const PetscInt **)idx, (const PetscScalar **)v)); 813ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 82ba8c8a56SBarry Smith } 83ba8c8a56SBarry Smith 8466976f2fSJacob Faibussowitsch static PetscErrorCode MatRestoreRow_MPIDense(Mat A, PetscInt row, PetscInt *nz, PetscInt **idx, PetscScalar **v) 85d71ae5a4SJacob Faibussowitsch { 86637a0070SStefano Zampini Mat_MPIDense *mat = (Mat_MPIDense *)A->data; 87637a0070SStefano Zampini PetscInt lrow, rstart = A->rmap->rstart, rend = A->rmap->rend; 88ba8c8a56SBarry Smith 89ba8c8a56SBarry Smith PetscFunctionBegin; 90aed4548fSBarry Smith PetscCheck(row >= rstart && row < rend, PETSC_COMM_SELF, PETSC_ERR_SUP, "only local rows"); 91637a0070SStefano Zampini lrow = row - rstart; 929566063dSJacob Faibussowitsch PetscCall(MatRestoreRow(mat->A, lrow, nz, (const PetscInt **)idx, (const PetscScalar **)v)); 933ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 94ba8c8a56SBarry Smith } 95ba8c8a56SBarry Smith 9666976f2fSJacob Faibussowitsch static PetscErrorCode MatGetDiagonalBlock_MPIDense(Mat A, Mat *a) 97d71ae5a4SJacob Faibussowitsch { 980de54da6SSatish Balay Mat_MPIDense *mdn = (Mat_MPIDense *)A->data; 99d0f46423SBarry Smith PetscInt m = A->rmap->n, rstart = A->rmap->rstart; 10087828ca2SBarry Smith PetscScalar *array; 1010de54da6SSatish Balay MPI_Comm comm; 1024b6ae80fSStefano Zampini PetscBool flg; 10311bd1e4dSLisandro Dalcin Mat B; 1040de54da6SSatish Balay 1050de54da6SSatish Balay PetscFunctionBegin; 1069566063dSJacob Faibussowitsch PetscCall(MatHasCongruentLayouts(A, &flg)); 10728b400f6SJacob Faibussowitsch PetscCheck(flg, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only square matrices supported."); 1089566063dSJacob Faibussowitsch PetscCall(PetscObjectQuery((PetscObject)A, "DiagonalBlock", (PetscObject *)&B)); 109616b8fbbSStefano Zampini if (!B) { /* This should use MatDenseGetSubMatrix (not create), but we would need a call like MatRestoreDiagonalBlock */ 110a51cd166SPierre Jolivet #if PetscDefined(HAVE_CUDA) 1119566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)mdn->A, MATSEQDENSECUDA, &flg)); 11228b400f6SJacob Faibussowitsch PetscCheck(!flg, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not coded for %s. Send an email to petsc-dev@mcs.anl.gov to request this feature", MATSEQDENSECUDA); 113a51cd166SPierre Jolivet #elif PetscDefined(HAVE_HIP) 11447d993e7Ssuyashtn PetscCall(PetscObjectTypeCompare((PetscObject)mdn->A, MATSEQDENSEHIP, &flg)); 11547d993e7Ssuyashtn PetscCheck(!flg, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not coded for %s. Send an email to petsc-dev@mcs.anl.gov to request this feature", MATSEQDENSEHIP); 11647d993e7Ssuyashtn #endif 117f4f49eeaSPierre Jolivet PetscCall(PetscObjectGetComm((PetscObject)mdn->A, &comm)); 1189566063dSJacob Faibussowitsch PetscCall(MatCreate(comm, &B)); 1199566063dSJacob Faibussowitsch PetscCall(MatSetSizes(B, m, m, m, m)); 1209566063dSJacob Faibussowitsch PetscCall(MatSetType(B, ((PetscObject)mdn->A)->type_name)); 1219566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(mdn->A, (const PetscScalar **)&array)); 1229566063dSJacob Faibussowitsch PetscCall(MatSeqDenseSetPreallocation(B, array + m * rstart)); 1239566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(mdn->A, (const PetscScalar **)&array)); 1249566063dSJacob Faibussowitsch PetscCall(PetscObjectCompose((PetscObject)A, "DiagonalBlock", (PetscObject)B)); 12511bd1e4dSLisandro Dalcin *a = B; 1269566063dSJacob Faibussowitsch PetscCall(MatDestroy(&B)); 1272205254eSKarl Rupp } else *a = B; 1283ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1290de54da6SSatish Balay } 1300de54da6SSatish Balay 13166976f2fSJacob Faibussowitsch static PetscErrorCode MatSetValues_MPIDense(Mat mat, PetscInt m, const PetscInt idxm[], PetscInt n, const PetscInt idxn[], const PetscScalar v[], InsertMode addv) 132d71ae5a4SJacob Faibussowitsch { 13339b7565bSBarry Smith Mat_MPIDense *A = (Mat_MPIDense *)mat->data; 134d0f46423SBarry Smith PetscInt i, j, rstart = mat->rmap->rstart, rend = mat->rmap->rend, row; 135ace3abfcSBarry Smith PetscBool roworiented = A->roworiented; 1368965ea79SLois Curfman McInnes 1373a40ed3dSBarry Smith PetscFunctionBegin; 1388965ea79SLois Curfman McInnes for (i = 0; i < m; i++) { 1395ef9f2a5SBarry Smith if (idxm[i] < 0) continue; 14008401ef6SPierre Jolivet PetscCheck(idxm[i] < mat->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large"); 1418965ea79SLois Curfman McInnes if (idxm[i] >= rstart && idxm[i] < rend) { 1428965ea79SLois Curfman McInnes row = idxm[i] - rstart; 14339b7565bSBarry Smith if (roworiented) { 144de0a9f35SPierre Jolivet PetscCall(MatSetValues(A->A, 1, &row, n, idxn, PetscSafePointerPlusOffset(v, i * n), addv)); 1453a40ed3dSBarry Smith } else { 1468965ea79SLois Curfman McInnes for (j = 0; j < n; j++) { 1475ef9f2a5SBarry Smith if (idxn[j] < 0) continue; 14808401ef6SPierre Jolivet PetscCheck(idxn[j] < mat->cmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Column too large"); 149de0a9f35SPierre Jolivet PetscCall(MatSetValues(A->A, 1, &row, 1, &idxn[j], PetscSafePointerPlusOffset(v, i + j * m), addv)); 15039b7565bSBarry Smith } 1518965ea79SLois Curfman McInnes } 1522205254eSKarl Rupp } else if (!A->donotstash) { 1535080c13bSMatthew G Knepley mat->assembled = PETSC_FALSE; 15439b7565bSBarry Smith if (roworiented) { 155de0a9f35SPierre Jolivet PetscCall(MatStashValuesRow_Private(&mat->stash, idxm[i], n, idxn, PetscSafePointerPlusOffset(v, i * n), PETSC_FALSE)); 156d36fbae8SSatish Balay } else { 157de0a9f35SPierre Jolivet PetscCall(MatStashValuesCol_Private(&mat->stash, idxm[i], n, idxn, PetscSafePointerPlusOffset(v, i), m, PETSC_FALSE)); 15839b7565bSBarry Smith } 159b49de8d1SLois Curfman McInnes } 160b49de8d1SLois Curfman McInnes } 1613ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 162b49de8d1SLois Curfman McInnes } 163b49de8d1SLois Curfman McInnes 16466976f2fSJacob Faibussowitsch static PetscErrorCode MatGetValues_MPIDense(Mat mat, PetscInt m, const PetscInt idxm[], PetscInt n, const PetscInt idxn[], PetscScalar v[]) 165d71ae5a4SJacob Faibussowitsch { 166b49de8d1SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data; 167d0f46423SBarry Smith PetscInt i, j, rstart = mat->rmap->rstart, rend = mat->rmap->rend, row; 168b49de8d1SLois Curfman McInnes 1693a40ed3dSBarry Smith PetscFunctionBegin; 170b49de8d1SLois Curfman McInnes for (i = 0; i < m; i++) { 17154c59aa7SJacob Faibussowitsch if (idxm[i] < 0) continue; /* negative row */ 17254c59aa7SJacob Faibussowitsch PetscCheck(idxm[i] < mat->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large"); 173b49de8d1SLois Curfman McInnes if (idxm[i] >= rstart && idxm[i] < rend) { 174b49de8d1SLois Curfman McInnes row = idxm[i] - rstart; 175b49de8d1SLois Curfman McInnes for (j = 0; j < n; j++) { 17654c59aa7SJacob Faibussowitsch if (idxn[j] < 0) continue; /* negative column */ 17754c59aa7SJacob Faibussowitsch PetscCheck(idxn[j] < mat->cmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Column too large"); 1789566063dSJacob Faibussowitsch PetscCall(MatGetValues(mdn->A, 1, &row, 1, &idxn[j], v + i * n + j)); 179b49de8d1SLois Curfman McInnes } 180e7e72b3dSBarry Smith } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "Only local values currently supported"); 1818965ea79SLois Curfman McInnes } 1823ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1838965ea79SLois Curfman McInnes } 1848965ea79SLois Curfman McInnes 185d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetLDA_MPIDense(Mat A, PetscInt *lda) 186d71ae5a4SJacob Faibussowitsch { 18749a6ff4bSBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data; 18849a6ff4bSBarry Smith 18949a6ff4bSBarry Smith PetscFunctionBegin; 1909566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, lda)); 1913ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 19249a6ff4bSBarry Smith } 19349a6ff4bSBarry Smith 194d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseSetLDA_MPIDense(Mat A, PetscInt lda) 195d71ae5a4SJacob Faibussowitsch { 196ad16ce7aSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data; 19747d993e7Ssuyashtn MatType mtype = MATSEQDENSE; 198ad16ce7aSStefano Zampini 199ad16ce7aSStefano Zampini PetscFunctionBegin; 20017359960SJose E. Roman if (!a->A) { 20128b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 2029566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(A->rmap)); 2039566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(A->cmap)); 2049566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF, &a->A)); 2059566063dSJacob Faibussowitsch PetscCall(MatSetSizes(a->A, A->rmap->n, A->cmap->N, A->rmap->n, A->cmap->N)); 206a51cd166SPierre Jolivet #if PetscDefined(HAVE_CUDA) 20747d993e7Ssuyashtn PetscBool iscuda; 2089566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSECUDA, &iscuda)); 20947d993e7Ssuyashtn if (iscuda) mtype = MATSEQDENSECUDA; 210a51cd166SPierre Jolivet #elif PetscDefined(HAVE_HIP) 21147d993e7Ssuyashtn PetscBool iship; 21247d993e7Ssuyashtn PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSEHIP, &iship)); 21347d993e7Ssuyashtn if (iship) mtype = MATSEQDENSEHIP; 21447d993e7Ssuyashtn #endif 21547d993e7Ssuyashtn PetscCall(MatSetType(a->A, mtype)); 21617359960SJose E. Roman } 2179566063dSJacob Faibussowitsch PetscCall(MatDenseSetLDA(a->A, lda)); 2183ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 219ad16ce7aSStefano Zampini } 220ad16ce7aSStefano Zampini 221d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetArray_MPIDense(Mat A, PetscScalar **array) 222d71ae5a4SJacob Faibussowitsch { 223ff14e315SSatish Balay Mat_MPIDense *a = (Mat_MPIDense *)A->data; 224ff14e315SSatish Balay 2253a40ed3dSBarry Smith PetscFunctionBegin; 22628b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 2279566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(a->A, array)); 2283ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 229ff14e315SSatish Balay } 230ff14e315SSatish Balay 231f1e99121SPierre Jolivet static PetscErrorCode MatDenseGetArrayRead_MPIDense(Mat A, PetscScalar **array) 232d71ae5a4SJacob Faibussowitsch { 2338572280aSBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data; 2348572280aSBarry Smith 2358572280aSBarry Smith PetscFunctionBegin; 23628b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 237f1e99121SPierre Jolivet PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)array)); 2383ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2398572280aSBarry Smith } 2408572280aSBarry Smith 241d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetArrayWrite_MPIDense(Mat A, PetscScalar **array) 242d71ae5a4SJacob Faibussowitsch { 2436947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data; 2446947451fSStefano Zampini 2456947451fSStefano Zampini PetscFunctionBegin; 24628b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 2479566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayWrite(a->A, array)); 2483ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2496947451fSStefano Zampini } 2506947451fSStefano Zampini 251d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDensePlaceArray_MPIDense(Mat A, const PetscScalar *array) 252d71ae5a4SJacob Faibussowitsch { 253d3042a70SBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data; 254d3042a70SBarry Smith 255d3042a70SBarry Smith PetscFunctionBegin; 25628b400f6SJacob Faibussowitsch PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 25728b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 2589566063dSJacob Faibussowitsch PetscCall(MatDensePlaceArray(a->A, array)); 2593ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 260d3042a70SBarry Smith } 261d3042a70SBarry Smith 262d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseResetArray_MPIDense(Mat A) 263d71ae5a4SJacob Faibussowitsch { 264d3042a70SBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data; 265d3042a70SBarry Smith 266d3042a70SBarry Smith PetscFunctionBegin; 26728b400f6SJacob Faibussowitsch PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 26828b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 2699566063dSJacob Faibussowitsch PetscCall(MatDenseResetArray(a->A)); 2703ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 271d3042a70SBarry Smith } 272d3042a70SBarry Smith 273d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseReplaceArray_MPIDense(Mat A, const PetscScalar *array) 274d71ae5a4SJacob Faibussowitsch { 275d5ea218eSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data; 276d5ea218eSStefano Zampini 277d5ea218eSStefano Zampini PetscFunctionBegin; 27828b400f6SJacob Faibussowitsch PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 27928b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 2809566063dSJacob Faibussowitsch PetscCall(MatDenseReplaceArray(a->A, array)); 2813ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 282d5ea218eSStefano Zampini } 283d5ea218eSStefano Zampini 284d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatCreateSubMatrix_MPIDense(Mat A, IS isrow, IS iscol, MatReuse scall, Mat *B) 285d71ae5a4SJacob Faibussowitsch { 286ca3fa75bSLois Curfman McInnes Mat_MPIDense *mat = (Mat_MPIDense *)A->data, *newmatd; 287637a0070SStefano Zampini PetscInt lda, i, j, rstart, rend, nrows, ncols, Ncols, nlrows, nlcols; 2885d0c19d7SBarry Smith const PetscInt *irow, *icol; 289637a0070SStefano Zampini const PetscScalar *v; 290637a0070SStefano Zampini PetscScalar *bv; 291ca3fa75bSLois Curfman McInnes Mat newmat; 2924aa3045dSJed Brown IS iscol_local; 29342a884f0SBarry Smith MPI_Comm comm_is, comm_mat; 294ca3fa75bSLois Curfman McInnes 295ca3fa75bSLois Curfman McInnes PetscFunctionBegin; 2969566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A, &comm_mat)); 2979566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)iscol, &comm_is)); 29808401ef6SPierre Jolivet PetscCheck(comm_mat == comm_is, PETSC_COMM_SELF, PETSC_ERR_ARG_NOTSAMECOMM, "IS communicator must match matrix communicator"); 29942a884f0SBarry Smith 3009566063dSJacob Faibussowitsch PetscCall(ISAllGather(iscol, &iscol_local)); 3019566063dSJacob Faibussowitsch PetscCall(ISGetIndices(isrow, &irow)); 3029566063dSJacob Faibussowitsch PetscCall(ISGetIndices(iscol_local, &icol)); 3039566063dSJacob Faibussowitsch PetscCall(ISGetLocalSize(isrow, &nrows)); 3049566063dSJacob Faibussowitsch PetscCall(ISGetLocalSize(iscol, &ncols)); 3059566063dSJacob Faibussowitsch PetscCall(ISGetSize(iscol, &Ncols)); /* global number of columns, size of iscol_local */ 306ca3fa75bSLois Curfman McInnes 307ca3fa75bSLois Curfman McInnes /* No parallel redistribution currently supported! Should really check each index set 308da81f932SPierre Jolivet to confirm that it is OK. ... Currently supports only submatrix same partitioning as 3097eba5e9cSLois Curfman McInnes original matrix! */ 310ca3fa75bSLois Curfman McInnes 3119566063dSJacob Faibussowitsch PetscCall(MatGetLocalSize(A, &nlrows, &nlcols)); 3129566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRange(A, &rstart, &rend)); 313ca3fa75bSLois Curfman McInnes 314ca3fa75bSLois Curfman McInnes /* Check submatrix call */ 315ca3fa75bSLois Curfman McInnes if (scall == MAT_REUSE_MATRIX) { 316e32f2f54SBarry Smith /* SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Reused submatrix wrong size"); */ 3177eba5e9cSLois Curfman McInnes /* Really need to test rows and column sizes! */ 318ca3fa75bSLois Curfman McInnes newmat = *B; 319ca3fa75bSLois Curfman McInnes } else { 320ca3fa75bSLois Curfman McInnes /* Create and fill new matrix */ 3219566063dSJacob Faibussowitsch PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &newmat)); 3229566063dSJacob Faibussowitsch PetscCall(MatSetSizes(newmat, nrows, ncols, PETSC_DECIDE, Ncols)); 3239566063dSJacob Faibussowitsch PetscCall(MatSetType(newmat, ((PetscObject)A)->type_name)); 3249566063dSJacob Faibussowitsch PetscCall(MatMPIDenseSetPreallocation(newmat, NULL)); 325ca3fa75bSLois Curfman McInnes } 326ca3fa75bSLois Curfman McInnes 327ca3fa75bSLois Curfman McInnes /* Now extract the data pointers and do the copy, column at a time */ 328ca3fa75bSLois Curfman McInnes newmatd = (Mat_MPIDense *)newmat->data; 3299566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(newmatd->A, &bv)); 3309566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(mat->A, &v)); 3319566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(mat->A, &lda)); 3324aa3045dSJed Brown for (i = 0; i < Ncols; i++) { 333637a0070SStefano Zampini const PetscScalar *av = v + lda * icol[i]; 334ad540459SPierre Jolivet for (j = 0; j < nrows; j++) *bv++ = av[irow[j] - rstart]; 335ca3fa75bSLois Curfman McInnes } 3369566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(mat->A, &v)); 3379566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(newmatd->A, &bv)); 338ca3fa75bSLois Curfman McInnes 339ca3fa75bSLois Curfman McInnes /* Assemble the matrices so that the correct flags are set */ 3409566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(newmat, MAT_FINAL_ASSEMBLY)); 3419566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(newmat, MAT_FINAL_ASSEMBLY)); 342ca3fa75bSLois Curfman McInnes 343ca3fa75bSLois Curfman McInnes /* Free work space */ 3449566063dSJacob Faibussowitsch PetscCall(ISRestoreIndices(isrow, &irow)); 3459566063dSJacob Faibussowitsch PetscCall(ISRestoreIndices(iscol_local, &icol)); 3469566063dSJacob Faibussowitsch PetscCall(ISDestroy(&iscol_local)); 347ca3fa75bSLois Curfman McInnes *B = newmat; 3483ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 349ca3fa75bSLois Curfman McInnes } 350ca3fa75bSLois Curfman McInnes 35166976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreArray_MPIDense(Mat A, PetscScalar **array) 352d71ae5a4SJacob Faibussowitsch { 35373a71a0fSBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data; 35473a71a0fSBarry Smith 3553a40ed3dSBarry Smith PetscFunctionBegin; 3569566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(a->A, array)); 3573ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 358ff14e315SSatish Balay } 359ff14e315SSatish Balay 360f1e99121SPierre Jolivet static PetscErrorCode MatDenseRestoreArrayRead_MPIDense(Mat A, PetscScalar **array) 361d71ae5a4SJacob Faibussowitsch { 3628572280aSBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data; 3638572280aSBarry Smith 3648572280aSBarry Smith PetscFunctionBegin; 365f1e99121SPierre Jolivet PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)array)); 3663ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 3678572280aSBarry Smith } 3688572280aSBarry Smith 36966976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreArrayWrite_MPIDense(Mat A, PetscScalar **array) 370d71ae5a4SJacob Faibussowitsch { 3716947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data; 3726947451fSStefano Zampini 3736947451fSStefano Zampini PetscFunctionBegin; 3749566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayWrite(a->A, array)); 3753ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 3766947451fSStefano Zampini } 3776947451fSStefano Zampini 37866976f2fSJacob Faibussowitsch static PetscErrorCode MatAssemblyBegin_MPIDense(Mat mat, MatAssemblyType mode) 379d71ae5a4SJacob Faibussowitsch { 38039ddd567SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data; 38113f74950SBarry Smith PetscInt nstash, reallocs; 3828965ea79SLois Curfman McInnes 3833a40ed3dSBarry Smith PetscFunctionBegin; 3843ba16761SJacob Faibussowitsch if (mdn->donotstash || mat->nooffprocentries) PetscFunctionReturn(PETSC_SUCCESS); 3858965ea79SLois Curfman McInnes 3869566063dSJacob Faibussowitsch PetscCall(MatStashScatterBegin_Private(mat, &mat->stash, mat->rmap->range)); 3879566063dSJacob Faibussowitsch PetscCall(MatStashGetInfo_Private(&mat->stash, &nstash, &reallocs)); 3889566063dSJacob Faibussowitsch PetscCall(PetscInfo(mdn->A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs)); 3893ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 3908965ea79SLois Curfman McInnes } 3918965ea79SLois Curfman McInnes 39266976f2fSJacob Faibussowitsch static PetscErrorCode MatAssemblyEnd_MPIDense(Mat mat, MatAssemblyType mode) 393d71ae5a4SJacob Faibussowitsch { 39439ddd567SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data; 39513f74950SBarry Smith PetscInt i, *row, *col, flg, j, rstart, ncols; 39613f74950SBarry Smith PetscMPIInt n; 39787828ca2SBarry Smith PetscScalar *val; 3988965ea79SLois Curfman McInnes 3993a40ed3dSBarry Smith PetscFunctionBegin; 400910cf402Sprj- if (!mdn->donotstash && !mat->nooffprocentries) { 4018965ea79SLois Curfman McInnes /* wait on receives */ 4027ef1d9bdSSatish Balay while (1) { 4039566063dSJacob Faibussowitsch PetscCall(MatStashScatterGetMesg_Private(&mat->stash, &n, &row, &col, &val, &flg)); 4047ef1d9bdSSatish Balay if (!flg) break; 4058965ea79SLois Curfman McInnes 4067ef1d9bdSSatish Balay for (i = 0; i < n;) { 4077ef1d9bdSSatish Balay /* Now identify the consecutive vals belonging to the same row */ 4082205254eSKarl Rupp for (j = i, rstart = row[j]; j < n; j++) { 4092205254eSKarl Rupp if (row[j] != rstart) break; 4102205254eSKarl Rupp } 4117ef1d9bdSSatish Balay if (j < n) ncols = j - i; 4127ef1d9bdSSatish Balay else ncols = n - i; 4137ef1d9bdSSatish Balay /* Now assemble all these values with a single function call */ 4149566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIDense(mat, 1, row + i, ncols, col + i, val + i, mat->insertmode)); 4157ef1d9bdSSatish Balay i = j; 4168965ea79SLois Curfman McInnes } 4177ef1d9bdSSatish Balay } 4189566063dSJacob Faibussowitsch PetscCall(MatStashScatterEnd_Private(&mat->stash)); 419910cf402Sprj- } 4208965ea79SLois Curfman McInnes 4219566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(mdn->A, mode)); 4229566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(mdn->A, mode)); 4233ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 4248965ea79SLois Curfman McInnes } 4258965ea79SLois Curfman McInnes 42666976f2fSJacob Faibussowitsch static PetscErrorCode MatZeroEntries_MPIDense(Mat A) 427d71ae5a4SJacob Faibussowitsch { 42839ddd567SLois Curfman McInnes Mat_MPIDense *l = (Mat_MPIDense *)A->data; 4293a40ed3dSBarry Smith 4303a40ed3dSBarry Smith PetscFunctionBegin; 4319566063dSJacob Faibussowitsch PetscCall(MatZeroEntries(l->A)); 4323ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 4338965ea79SLois Curfman McInnes } 4348965ea79SLois Curfman McInnes 43566976f2fSJacob Faibussowitsch static PetscErrorCode MatZeroRows_MPIDense(Mat A, PetscInt n, const PetscInt rows[], PetscScalar diag, Vec x, Vec b) 436d71ae5a4SJacob Faibussowitsch { 43739ddd567SLois Curfman McInnes Mat_MPIDense *l = (Mat_MPIDense *)A->data; 438637a0070SStefano Zampini PetscInt i, len, *lrows; 439637a0070SStefano Zampini 440637a0070SStefano Zampini PetscFunctionBegin; 441637a0070SStefano Zampini /* get locally owned rows */ 4429566063dSJacob Faibussowitsch PetscCall(PetscLayoutMapLocal(A->rmap, n, rows, &len, &lrows, NULL)); 443dd8e379bSPierre Jolivet /* fix right-hand side if needed */ 444637a0070SStefano Zampini if (x && b) { 44597b48c8fSBarry Smith const PetscScalar *xx; 44697b48c8fSBarry Smith PetscScalar *bb; 4478965ea79SLois Curfman McInnes 4489566063dSJacob Faibussowitsch PetscCall(VecGetArrayRead(x, &xx)); 4499566063dSJacob Faibussowitsch PetscCall(VecGetArrayWrite(b, &bb)); 450637a0070SStefano Zampini for (i = 0; i < len; ++i) bb[lrows[i]] = diag * xx[lrows[i]]; 4519566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayRead(x, &xx)); 4529566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayWrite(b, &bb)); 45397b48c8fSBarry Smith } 4549566063dSJacob Faibussowitsch PetscCall(MatZeroRows(l->A, len, lrows, 0.0, NULL, NULL)); 455e2eb51b1SBarry Smith if (diag != 0.0) { 456637a0070SStefano Zampini Vec d; 457b9679d65SBarry Smith 4589566063dSJacob Faibussowitsch PetscCall(MatCreateVecs(A, NULL, &d)); 4599566063dSJacob Faibussowitsch PetscCall(VecSet(d, diag)); 4609566063dSJacob Faibussowitsch PetscCall(MatDiagonalSet(A, d, INSERT_VALUES)); 4619566063dSJacob Faibussowitsch PetscCall(VecDestroy(&d)); 462b9679d65SBarry Smith } 4639566063dSJacob Faibussowitsch PetscCall(PetscFree(lrows)); 4643ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 4658965ea79SLois Curfman McInnes } 4668965ea79SLois Curfman McInnes 4670be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMult_SeqDense(Mat, Vec, Vec); 4680be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMultAdd_SeqDense(Mat, Vec, Vec, Vec); 4690be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMultTranspose_SeqDense(Mat, Vec, Vec); 4700be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMultTransposeAdd_SeqDense(Mat, Vec, Vec, Vec); 4710be0d8bdSHansol Suh 47266976f2fSJacob Faibussowitsch static PetscErrorCode MatMult_MPIDense(Mat mat, Vec xx, Vec yy) 473d71ae5a4SJacob Faibussowitsch { 47439ddd567SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data; 475637a0070SStefano Zampini const PetscScalar *ax; 476637a0070SStefano Zampini PetscScalar *ay; 477d0295fc0SJunchao Zhang PetscMemType axmtype, aymtype; 478c456f294SBarry Smith 4793a40ed3dSBarry Smith PetscFunctionBegin; 4806f08ad05SPierre Jolivet if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat)); 4819566063dSJacob Faibussowitsch PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype)); 4829566063dSJacob Faibussowitsch PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype)); 4839566063dSJacob Faibussowitsch PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE)); 4849566063dSJacob Faibussowitsch PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE)); 4859566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay)); 4869566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayReadAndMemType(xx, &ax)); 4879566063dSJacob Faibussowitsch PetscCall((*mdn->A->ops->mult)(mdn->A, mdn->lvec, yy)); 4883ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 4898965ea79SLois Curfman McInnes } 4908965ea79SLois Curfman McInnes 4910be0d8bdSHansol Suh static PetscErrorCode MatMultAddColumnRange_MPIDense(Mat mat, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end) 4920be0d8bdSHansol Suh { 4930be0d8bdSHansol Suh Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data; 4940be0d8bdSHansol Suh const PetscScalar *ax; 4950be0d8bdSHansol Suh PetscScalar *ay; 4960be0d8bdSHansol Suh PetscMemType axmtype, aymtype; 4970be0d8bdSHansol Suh 4980be0d8bdSHansol Suh PetscFunctionBegin; 4990be0d8bdSHansol Suh if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat)); 5000be0d8bdSHansol Suh PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype)); 5010be0d8bdSHansol Suh PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype)); 5020be0d8bdSHansol Suh PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE)); 5030be0d8bdSHansol Suh PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE)); 5040be0d8bdSHansol Suh PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay)); 5050be0d8bdSHansol Suh PetscCall(VecRestoreArrayReadAndMemType(xx, &ax)); 5060be0d8bdSHansol Suh PetscUseMethod(mdn->A, "MatMultAddColumnRange_C", (Mat, Vec, Vec, Vec, PetscInt, PetscInt), (mdn->A, mdn->lvec, yy, zz, c_start, c_end)); 5070be0d8bdSHansol Suh PetscFunctionReturn(PETSC_SUCCESS); 5080be0d8bdSHansol Suh } 5090be0d8bdSHansol Suh 51066976f2fSJacob Faibussowitsch static PetscErrorCode MatMultAdd_MPIDense(Mat mat, Vec xx, Vec yy, Vec zz) 511d71ae5a4SJacob Faibussowitsch { 51239ddd567SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data; 513637a0070SStefano Zampini const PetscScalar *ax; 514637a0070SStefano Zampini PetscScalar *ay; 515d0295fc0SJunchao Zhang PetscMemType axmtype, aymtype; 516c456f294SBarry Smith 5173a40ed3dSBarry Smith PetscFunctionBegin; 5186f08ad05SPierre Jolivet if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat)); 5199566063dSJacob Faibussowitsch PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype)); 5209566063dSJacob Faibussowitsch PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype)); 5219566063dSJacob Faibussowitsch PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE)); 5229566063dSJacob Faibussowitsch PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE)); 5239566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay)); 5249566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayReadAndMemType(xx, &ax)); 5259566063dSJacob Faibussowitsch PetscCall((*mdn->A->ops->multadd)(mdn->A, mdn->lvec, yy, zz)); 5263ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 5278965ea79SLois Curfman McInnes } 5288965ea79SLois Curfman McInnes 5290be0d8bdSHansol Suh static PetscErrorCode MatMultHermitianTransposeColumnRange_MPIDense(Mat A, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end) 5300be0d8bdSHansol Suh { 5310be0d8bdSHansol Suh Mat_MPIDense *a = (Mat_MPIDense *)A->data; 5320be0d8bdSHansol Suh const PetscScalar *ax; 5330be0d8bdSHansol Suh PetscScalar *ay; 5340be0d8bdSHansol Suh PetscMemType axmtype, aymtype; 5350be0d8bdSHansol Suh 5360be0d8bdSHansol Suh PetscFunctionBegin; 5370be0d8bdSHansol Suh if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A)); 5380be0d8bdSHansol Suh PetscCall(VecSet(yy, 0.0)); 5390be0d8bdSHansol Suh PetscUseMethod(a->A, "MatMultHermitianTransposeColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (a->A, xx, a->lvec, c_start, c_end)); 5400be0d8bdSHansol Suh PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype)); 5410be0d8bdSHansol Suh PetscCall(VecGetArrayAndMemType(yy, &ay, &aymtype)); 5420be0d8bdSHansol Suh PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM)); 5430be0d8bdSHansol Suh PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM)); 5440be0d8bdSHansol Suh PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax)); 5450be0d8bdSHansol Suh PetscCall(VecRestoreArrayAndMemType(yy, &ay)); 5460be0d8bdSHansol Suh PetscFunctionReturn(PETSC_SUCCESS); 5470be0d8bdSHansol Suh } 5480be0d8bdSHansol Suh 549459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTransposeKernel_MPIDense(Mat A, Vec xx, Vec yy, PetscBool herm) 550d71ae5a4SJacob Faibussowitsch { 551096963f5SLois Curfman McInnes Mat_MPIDense *a = (Mat_MPIDense *)A->data; 552637a0070SStefano Zampini const PetscScalar *ax; 553637a0070SStefano Zampini PetscScalar *ay; 554d0295fc0SJunchao Zhang PetscMemType axmtype, aymtype; 5554d86920dSPierre Jolivet 5563a40ed3dSBarry Smith PetscFunctionBegin; 5576f08ad05SPierre Jolivet if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A)); 5589566063dSJacob Faibussowitsch PetscCall(VecSet(yy, 0.0)); 559459e8d23SBlanca Mellado Pinto if (herm) PetscCall((*a->A->ops->multhermitiantranspose)(a->A, xx, a->lvec)); 560459e8d23SBlanca Mellado Pinto else PetscCall((*a->A->ops->multtranspose)(a->A, xx, a->lvec)); 5619566063dSJacob Faibussowitsch PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype)); 5629566063dSJacob Faibussowitsch PetscCall(VecGetArrayAndMemType(yy, &ay, &aymtype)); 5639566063dSJacob Faibussowitsch PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM)); 5649566063dSJacob Faibussowitsch PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM)); 5659566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax)); 5669566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayAndMemType(yy, &ay)); 5673ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 568096963f5SLois Curfman McInnes } 569096963f5SLois Curfman McInnes 5700be0d8bdSHansol Suh static PetscErrorCode MatMultHermitianTransposeAddColumnRange_MPIDense(Mat A, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end) 5710be0d8bdSHansol Suh { 5720be0d8bdSHansol Suh Mat_MPIDense *a = (Mat_MPIDense *)A->data; 5730be0d8bdSHansol Suh const PetscScalar *ax; 5740be0d8bdSHansol Suh PetscScalar *ay; 5750be0d8bdSHansol Suh PetscMemType axmtype, aymtype; 5760be0d8bdSHansol Suh 5770be0d8bdSHansol Suh PetscFunctionBegin; 5780be0d8bdSHansol Suh if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A)); 5790be0d8bdSHansol Suh PetscCall(VecCopy(yy, zz)); 5800be0d8bdSHansol Suh PetscMPIInt rank; 5810be0d8bdSHansol Suh PetscCallMPI(MPI_Comm_rank(MPI_COMM_WORLD, &rank)); 5820be0d8bdSHansol Suh PetscUseMethod(a->A, "MatMultHermitianTransposeColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (a->A, xx, a->lvec, c_start, c_end)); 5830be0d8bdSHansol Suh PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype)); 5840be0d8bdSHansol Suh PetscCall(VecGetArrayAndMemType(zz, &ay, &aymtype)); 5850be0d8bdSHansol Suh PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM)); 5860be0d8bdSHansol Suh PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM)); 5870be0d8bdSHansol Suh PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax)); 5880be0d8bdSHansol Suh PetscCall(VecRestoreArrayAndMemType(zz, &ay)); 5890be0d8bdSHansol Suh PetscFunctionReturn(PETSC_SUCCESS); 5900be0d8bdSHansol Suh } 5910be0d8bdSHansol Suh 592459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTransposeAddKernel_MPIDense(Mat A, Vec xx, Vec yy, Vec zz, PetscBool herm) 593d71ae5a4SJacob Faibussowitsch { 594096963f5SLois Curfman McInnes Mat_MPIDense *a = (Mat_MPIDense *)A->data; 595637a0070SStefano Zampini const PetscScalar *ax; 596637a0070SStefano Zampini PetscScalar *ay; 597d0295fc0SJunchao Zhang PetscMemType axmtype, aymtype; 598096963f5SLois Curfman McInnes 5993a40ed3dSBarry Smith PetscFunctionBegin; 6006f08ad05SPierre Jolivet if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A)); 6019566063dSJacob Faibussowitsch PetscCall(VecCopy(yy, zz)); 602459e8d23SBlanca Mellado Pinto if (herm) PetscCall((*a->A->ops->multhermitiantranspose)(a->A, xx, a->lvec)); 603459e8d23SBlanca Mellado Pinto else PetscCall((*a->A->ops->multtranspose)(a->A, xx, a->lvec)); 6049566063dSJacob Faibussowitsch PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype)); 6059566063dSJacob Faibussowitsch PetscCall(VecGetArrayAndMemType(zz, &ay, &aymtype)); 6069566063dSJacob Faibussowitsch PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM)); 6079566063dSJacob Faibussowitsch PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM)); 6089566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax)); 6099566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayAndMemType(zz, &ay)); 6103ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 611096963f5SLois Curfman McInnes } 612096963f5SLois Curfman McInnes 613459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTranspose_MPIDense(Mat A, Vec xx, Vec yy) 614459e8d23SBlanca Mellado Pinto { 615459e8d23SBlanca Mellado Pinto PetscFunctionBegin; 616459e8d23SBlanca Mellado Pinto PetscCall(MatMultTransposeKernel_MPIDense(A, xx, yy, PETSC_FALSE)); 617459e8d23SBlanca Mellado Pinto PetscFunctionReturn(PETSC_SUCCESS); 618459e8d23SBlanca Mellado Pinto } 619459e8d23SBlanca Mellado Pinto 620459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTransposeAdd_MPIDense(Mat A, Vec xx, Vec yy, Vec zz) 621459e8d23SBlanca Mellado Pinto { 622459e8d23SBlanca Mellado Pinto PetscFunctionBegin; 623459e8d23SBlanca Mellado Pinto PetscCall(MatMultTransposeAddKernel_MPIDense(A, xx, yy, zz, PETSC_FALSE)); 624459e8d23SBlanca Mellado Pinto PetscFunctionReturn(PETSC_SUCCESS); 625459e8d23SBlanca Mellado Pinto } 626459e8d23SBlanca Mellado Pinto 627459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultHermitianTranspose_MPIDense(Mat A, Vec xx, Vec yy) 628459e8d23SBlanca Mellado Pinto { 629459e8d23SBlanca Mellado Pinto PetscFunctionBegin; 630459e8d23SBlanca Mellado Pinto PetscCall(MatMultTransposeKernel_MPIDense(A, xx, yy, PETSC_TRUE)); 631459e8d23SBlanca Mellado Pinto PetscFunctionReturn(PETSC_SUCCESS); 632459e8d23SBlanca Mellado Pinto } 633459e8d23SBlanca Mellado Pinto 634459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultHermitianTransposeAdd_MPIDense(Mat A, Vec xx, Vec yy, Vec zz) 635459e8d23SBlanca Mellado Pinto { 636459e8d23SBlanca Mellado Pinto PetscFunctionBegin; 637459e8d23SBlanca Mellado Pinto PetscCall(MatMultTransposeAddKernel_MPIDense(A, xx, yy, zz, PETSC_TRUE)); 638459e8d23SBlanca Mellado Pinto PetscFunctionReturn(PETSC_SUCCESS); 639459e8d23SBlanca Mellado Pinto } 640459e8d23SBlanca Mellado Pinto 641d71ae5a4SJacob Faibussowitsch PetscErrorCode MatGetDiagonal_MPIDense(Mat A, Vec v) 642d71ae5a4SJacob Faibussowitsch { 64339ddd567SLois Curfman McInnes Mat_MPIDense *a = (Mat_MPIDense *)A->data; 64479560b15SJacob Faibussowitsch PetscInt lda, len, i, nl, ng, m = A->rmap->n, radd; 64579560b15SJacob Faibussowitsch PetscScalar *x; 646637a0070SStefano Zampini const PetscScalar *av; 647ed3cc1f0SBarry Smith 6483a40ed3dSBarry Smith PetscFunctionBegin; 6499566063dSJacob Faibussowitsch PetscCall(VecGetArray(v, &x)); 65079560b15SJacob Faibussowitsch PetscCall(VecGetSize(v, &ng)); 65179560b15SJacob Faibussowitsch PetscCheck(ng == A->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Nonconforming mat and vec"); 65279560b15SJacob Faibussowitsch PetscCall(VecGetLocalSize(v, &nl)); 653d0f46423SBarry Smith len = PetscMin(a->A->rmap->n, a->A->cmap->n); 654d0f46423SBarry Smith radd = A->rmap->rstart * m; 6559566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(a->A, &av)); 6569566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &lda)); 657ad540459SPierre Jolivet for (i = 0; i < len; i++) x[i] = av[radd + i * lda + i]; 6589566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(a->A, &av)); 6590be0d8bdSHansol Suh if (nl - i > 0) PetscCall(PetscArrayzero(x + i, nl - i)); 6609566063dSJacob Faibussowitsch PetscCall(VecRestoreArray(v, &x)); 6613ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 6628965ea79SLois Curfman McInnes } 6638965ea79SLois Curfman McInnes 66466976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MPIDense(Mat mat) 665d71ae5a4SJacob Faibussowitsch { 6663501a2bdSLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data; 667ed3cc1f0SBarry Smith 6683a40ed3dSBarry Smith PetscFunctionBegin; 6693ba16761SJacob Faibussowitsch PetscCall(PetscLogObjectState((PetscObject)mat, "Rows=%" PetscInt_FMT ", Cols=%" PetscInt_FMT, mat->rmap->N, mat->cmap->N)); 6709566063dSJacob Faibussowitsch PetscCall(MatStashDestroy_Private(&mat->stash)); 67128b400f6SJacob Faibussowitsch PetscCheck(!mdn->vecinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 67228b400f6SJacob Faibussowitsch PetscCheck(!mdn->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 6739566063dSJacob Faibussowitsch PetscCall(MatDestroy(&mdn->A)); 6749566063dSJacob Faibussowitsch PetscCall(VecDestroy(&mdn->lvec)); 6759566063dSJacob Faibussowitsch PetscCall(PetscSFDestroy(&mdn->Mvctx)); 6769566063dSJacob Faibussowitsch PetscCall(VecDestroy(&mdn->cvec)); 6779566063dSJacob Faibussowitsch PetscCall(MatDestroy(&mdn->cmat)); 67801b82886SBarry Smith 6799566063dSJacob Faibussowitsch PetscCall(PetscFree(mat->data)); 6809566063dSJacob Faibussowitsch PetscCall(PetscObjectChangeTypeName((PetscObject)mat, NULL)); 6818baccfbdSHong Zhang 6829566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", NULL)); 6839566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", NULL)); 6849566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", NULL)); 6859566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", NULL)); 6869566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", NULL)); 6879566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", NULL)); 6889566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", NULL)); 6899566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", NULL)); 6909566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", NULL)); 6919566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", NULL)); 6929566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", NULL)); 6939566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", NULL)); 6949566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", NULL)); 6958baccfbdSHong Zhang #if defined(PETSC_HAVE_ELEMENTAL) 6969566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", NULL)); 6978baccfbdSHong Zhang #endif 698d24d4204SJose E. Roman #if defined(PETSC_HAVE_SCALAPACK) 6999566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", NULL)); 700d24d4204SJose E. Roman #endif 7019566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", NULL)); 7029566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", NULL)); 7039566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", NULL)); 7046718818eSStefano Zampini #if defined(PETSC_HAVE_CUDA) 7059566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", NULL)); 7069566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", NULL)); 7072e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", NULL)); 7082e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidensecuda_mpidense_C", NULL)); 7092e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidensecuda_C", NULL)); 7102e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidensecuda_C", NULL)); 7112e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensecuda_mpiaij_C", NULL)); 7122e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensecuda_mpiaijcusparse_C", NULL)); 7132e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArray_C", NULL)); 7142e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArrayRead_C", NULL)); 7152e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArrayWrite_C", NULL)); 7162e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArray_C", NULL)); 7172e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArrayRead_C", NULL)); 7182e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArrayWrite_C", NULL)); 7192e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAPlaceArray_C", NULL)); 7202e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAResetArray_C", NULL)); 7212e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAReplaceArray_C", NULL)); 7223d9668e3SJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDASetPreallocation_C", NULL)); 7236718818eSStefano Zampini #endif 72447d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP) 72547d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", NULL)); 72647d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", NULL)); 72747d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", NULL)); 72847d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidensehip_mpidense_C", NULL)); 72947d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidensehip_C", NULL)); 73047d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidensehip_C", NULL)); 73147d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensehip_mpiaij_C", NULL)); 73247d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensehip_mpiaijhipsparse_C", NULL)); 73347d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArray_C", NULL)); 73447d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArrayRead_C", NULL)); 73547d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArrayWrite_C", NULL)); 73647d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArray_C", NULL)); 73747d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArrayRead_C", NULL)); 73847d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArrayWrite_C", NULL)); 73947d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPPlaceArray_C", NULL)); 74047d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPResetArray_C", NULL)); 74147d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPReplaceArray_C", NULL)); 7423d9668e3SJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPSetPreallocation_C", NULL)); 74347d993e7Ssuyashtn #endif 7449566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", NULL)); 7459566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", NULL)); 7469566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", NULL)); 7479566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", NULL)); 7489566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", NULL)); 7499566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", NULL)); 7509566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", NULL)); 7519566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", NULL)); 7529566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", NULL)); 7539566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", NULL)); 7540be0d8bdSHansol Suh PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", NULL)); 7550be0d8bdSHansol Suh PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", NULL)); 7560be0d8bdSHansol Suh PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", NULL)); 757f8103e6bSStefano Zampini 758f8103e6bSStefano Zampini PetscCall(PetscObjectCompose((PetscObject)mat, "DiagonalBlock", NULL)); 7593ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 7608965ea79SLois Curfman McInnes } 76139ddd567SLois Curfman McInnes 7629804daf3SBarry Smith #include <petscdraw.h> 763d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatView_MPIDense_ASCIIorDraworSocket(Mat mat, PetscViewer viewer) 764d71ae5a4SJacob Faibussowitsch { 76539ddd567SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data; 766616b8fbbSStefano Zampini PetscMPIInt rank; 76719fd82e9SBarry Smith PetscViewerType vtype; 768ace3abfcSBarry Smith PetscBool iascii, isdraw; 769b0a32e0cSBarry Smith PetscViewer sviewer; 770f3ef73ceSBarry Smith PetscViewerFormat format; 7718965ea79SLois Curfman McInnes 7723a40ed3dSBarry Smith PetscFunctionBegin; 7739566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mat), &rank)); 7749566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii)); 7759566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw)); 77632077d6dSBarry Smith if (iascii) { 7779566063dSJacob Faibussowitsch PetscCall(PetscViewerGetType(viewer, &vtype)); 7789566063dSJacob Faibussowitsch PetscCall(PetscViewerGetFormat(viewer, &format)); 779456192e2SBarry Smith if (format == PETSC_VIEWER_ASCII_INFO_DETAIL) { 7804e220ebcSLois Curfman McInnes MatInfo info; 7819566063dSJacob Faibussowitsch PetscCall(MatGetInfo(mat, MAT_LOCAL, &info)); 7829566063dSJacob Faibussowitsch PetscCall(PetscViewerASCIIPushSynchronized(viewer)); 7839371c9d4SSatish Balay PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, " [%d] local rows %" PetscInt_FMT " nz %" PetscInt_FMT " nz alloced %" PetscInt_FMT " mem %" PetscInt_FMT " \n", rank, mat->rmap->n, (PetscInt)info.nz_used, (PetscInt)info.nz_allocated, 7849371c9d4SSatish Balay (PetscInt)info.memory)); 7859566063dSJacob Faibussowitsch PetscCall(PetscViewerFlush(viewer)); 7869566063dSJacob Faibussowitsch PetscCall(PetscViewerASCIIPopSynchronized(viewer)); 7876f08ad05SPierre Jolivet if (mdn->Mvctx) PetscCall(PetscSFView(mdn->Mvctx, viewer)); 7883ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 789fb9695e5SSatish Balay } else if (format == PETSC_VIEWER_ASCII_INFO) { 7903ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 7918965ea79SLois Curfman McInnes } 792f1af5d2fSBarry Smith } else if (isdraw) { 793b0a32e0cSBarry Smith PetscDraw draw; 794ace3abfcSBarry Smith PetscBool isnull; 795f1af5d2fSBarry Smith 7969566063dSJacob Faibussowitsch PetscCall(PetscViewerDrawGetDraw(viewer, 0, &draw)); 7979566063dSJacob Faibussowitsch PetscCall(PetscDrawIsNull(draw, &isnull)); 7983ba16761SJacob Faibussowitsch if (isnull) PetscFunctionReturn(PETSC_SUCCESS); 799f1af5d2fSBarry Smith } 80077ed5343SBarry Smith 8017da1fb6eSBarry Smith { 8028965ea79SLois Curfman McInnes /* assemble the entire matrix onto first processor. */ 8038965ea79SLois Curfman McInnes Mat A; 804d0f46423SBarry Smith PetscInt M = mat->rmap->N, N = mat->cmap->N, m, row, i, nz; 805ba8c8a56SBarry Smith PetscInt *cols; 806ba8c8a56SBarry Smith PetscScalar *vals; 8078965ea79SLois Curfman McInnes 8089566063dSJacob Faibussowitsch PetscCall(MatCreate(PetscObjectComm((PetscObject)mat), &A)); 809dd400576SPatrick Sanan if (rank == 0) { 8109566063dSJacob Faibussowitsch PetscCall(MatSetSizes(A, M, N, M, N)); 8113a40ed3dSBarry Smith } else { 8129566063dSJacob Faibussowitsch PetscCall(MatSetSizes(A, 0, 0, M, N)); 8138965ea79SLois Curfman McInnes } 8147adad957SLisandro Dalcin /* Since this is a temporary matrix, MATMPIDENSE instead of ((PetscObject)A)->type_name here is probably acceptable. */ 8159566063dSJacob Faibussowitsch PetscCall(MatSetType(A, MATMPIDENSE)); 8169566063dSJacob Faibussowitsch PetscCall(MatMPIDenseSetPreallocation(A, NULL)); 8178965ea79SLois Curfman McInnes 81839ddd567SLois Curfman McInnes /* Copy the matrix ... This isn't the most efficient means, 81939ddd567SLois Curfman McInnes but it's quick for now */ 82051022da4SBarry Smith A->insertmode = INSERT_VALUES; 8212205254eSKarl Rupp 8222205254eSKarl Rupp row = mat->rmap->rstart; 8232205254eSKarl Rupp m = mdn->A->rmap->n; 8248965ea79SLois Curfman McInnes for (i = 0; i < m; i++) { 8259566063dSJacob Faibussowitsch PetscCall(MatGetRow_MPIDense(mat, row, &nz, &cols, &vals)); 8269566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIDense(A, 1, &row, nz, cols, vals, INSERT_VALUES)); 8279566063dSJacob Faibussowitsch PetscCall(MatRestoreRow_MPIDense(mat, row, &nz, &cols, &vals)); 82839ddd567SLois Curfman McInnes row++; 8298965ea79SLois Curfman McInnes } 8308965ea79SLois Curfman McInnes 8319566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 8329566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 8339566063dSJacob Faibussowitsch PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sviewer)); 834dd400576SPatrick Sanan if (rank == 0) { 835f4f49eeaSPierre Jolivet PetscCall(PetscObjectSetName((PetscObject)((Mat_MPIDense *)A->data)->A, ((PetscObject)mat)->name)); 836f4f49eeaSPierre Jolivet PetscCall(MatView_SeqDense(((Mat_MPIDense *)A->data)->A, sviewer)); 8378965ea79SLois Curfman McInnes } 8389566063dSJacob Faibussowitsch PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sviewer)); 8399566063dSJacob Faibussowitsch PetscCall(MatDestroy(&A)); 8408965ea79SLois Curfman McInnes } 8413ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 8428965ea79SLois Curfman McInnes } 8438965ea79SLois Curfman McInnes 84466976f2fSJacob Faibussowitsch static PetscErrorCode MatView_MPIDense(Mat mat, PetscViewer viewer) 845d71ae5a4SJacob Faibussowitsch { 846ace3abfcSBarry Smith PetscBool iascii, isbinary, isdraw, issocket; 8478965ea79SLois Curfman McInnes 848433994e6SBarry Smith PetscFunctionBegin; 8499566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii)); 8509566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary)); 8519566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERSOCKET, &issocket)); 8529566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw)); 8530f5bd95cSBarry Smith 85432077d6dSBarry Smith if (iascii || issocket || isdraw) { 8559566063dSJacob Faibussowitsch PetscCall(MatView_MPIDense_ASCIIorDraworSocket(mat, viewer)); 8561baa6e33SBarry Smith } else if (isbinary) PetscCall(MatView_Dense_Binary(mat, viewer)); 8573ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 8588965ea79SLois Curfman McInnes } 8598965ea79SLois Curfman McInnes 86066976f2fSJacob Faibussowitsch static PetscErrorCode MatGetInfo_MPIDense(Mat A, MatInfoType flag, MatInfo *info) 861d71ae5a4SJacob Faibussowitsch { 8623501a2bdSLois Curfman McInnes Mat_MPIDense *mat = (Mat_MPIDense *)A->data; 8633501a2bdSLois Curfman McInnes Mat mdn = mat->A; 8643966268fSBarry Smith PetscLogDouble isend[5], irecv[5]; 8658965ea79SLois Curfman McInnes 8663a40ed3dSBarry Smith PetscFunctionBegin; 8674e220ebcSLois Curfman McInnes info->block_size = 1.0; 8682205254eSKarl Rupp 8699566063dSJacob Faibussowitsch PetscCall(MatGetInfo(mdn, MAT_LOCAL, info)); 8702205254eSKarl Rupp 8719371c9d4SSatish Balay isend[0] = info->nz_used; 8729371c9d4SSatish Balay isend[1] = info->nz_allocated; 8739371c9d4SSatish Balay isend[2] = info->nz_unneeded; 8749371c9d4SSatish Balay isend[3] = info->memory; 8759371c9d4SSatish Balay isend[4] = info->mallocs; 8768965ea79SLois Curfman McInnes if (flag == MAT_LOCAL) { 8774e220ebcSLois Curfman McInnes info->nz_used = isend[0]; 8784e220ebcSLois Curfman McInnes info->nz_allocated = isend[1]; 8794e220ebcSLois Curfman McInnes info->nz_unneeded = isend[2]; 8804e220ebcSLois Curfman McInnes info->memory = isend[3]; 8814e220ebcSLois Curfman McInnes info->mallocs = isend[4]; 8828965ea79SLois Curfman McInnes } else if (flag == MAT_GLOBAL_MAX) { 8831c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(isend, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_MAX, PetscObjectComm((PetscObject)A))); 8842205254eSKarl Rupp 8854e220ebcSLois Curfman McInnes info->nz_used = irecv[0]; 8864e220ebcSLois Curfman McInnes info->nz_allocated = irecv[1]; 8874e220ebcSLois Curfman McInnes info->nz_unneeded = irecv[2]; 8884e220ebcSLois Curfman McInnes info->memory = irecv[3]; 8894e220ebcSLois Curfman McInnes info->mallocs = irecv[4]; 8908965ea79SLois Curfman McInnes } else if (flag == MAT_GLOBAL_SUM) { 8911c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(isend, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_SUM, PetscObjectComm((PetscObject)A))); 8922205254eSKarl Rupp 8934e220ebcSLois Curfman McInnes info->nz_used = irecv[0]; 8944e220ebcSLois Curfman McInnes info->nz_allocated = irecv[1]; 8954e220ebcSLois Curfman McInnes info->nz_unneeded = irecv[2]; 8964e220ebcSLois Curfman McInnes info->memory = irecv[3]; 8974e220ebcSLois Curfman McInnes info->mallocs = irecv[4]; 8988965ea79SLois Curfman McInnes } 8994e220ebcSLois Curfman McInnes info->fill_ratio_given = 0; /* no parallel LU/ILU/Cholesky */ 9004e220ebcSLois Curfman McInnes info->fill_ratio_needed = 0; 9014e220ebcSLois Curfman McInnes info->factor_mallocs = 0; 9023ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 9038965ea79SLois Curfman McInnes } 9048965ea79SLois Curfman McInnes 90566976f2fSJacob Faibussowitsch static PetscErrorCode MatSetOption_MPIDense(Mat A, MatOption op, PetscBool flg) 906d71ae5a4SJacob Faibussowitsch { 90739ddd567SLois Curfman McInnes Mat_MPIDense *a = (Mat_MPIDense *)A->data; 9088965ea79SLois Curfman McInnes 9093a40ed3dSBarry Smith PetscFunctionBegin; 91012c028f9SKris Buschelman switch (op) { 911512a5fc5SBarry Smith case MAT_NEW_NONZERO_LOCATIONS: 91212c028f9SKris Buschelman case MAT_NEW_NONZERO_LOCATION_ERR: 91312c028f9SKris Buschelman case MAT_NEW_NONZERO_ALLOCATION_ERR: 91443674050SBarry Smith MatCheckPreallocated(A, 1); 9159566063dSJacob Faibussowitsch PetscCall(MatSetOption(a->A, op, flg)); 91612c028f9SKris Buschelman break; 91712c028f9SKris Buschelman case MAT_ROW_ORIENTED: 91843674050SBarry Smith MatCheckPreallocated(A, 1); 9194e0d8c25SBarry Smith a->roworiented = flg; 9209566063dSJacob Faibussowitsch PetscCall(MatSetOption(a->A, op, flg)); 92112c028f9SKris Buschelman break; 9228c78258cSHong Zhang case MAT_FORCE_DIAGONAL_ENTRIES: 92313fa8e87SLisandro Dalcin case MAT_KEEP_NONZERO_PATTERN: 92412c028f9SKris Buschelman case MAT_USE_HASH_TABLE: 925d71ae5a4SJacob Faibussowitsch case MAT_SORTED_FULL: 926d71ae5a4SJacob Faibussowitsch PetscCall(PetscInfo(A, "Option %s ignored\n", MatOptions[op])); 927d71ae5a4SJacob Faibussowitsch break; 928d71ae5a4SJacob Faibussowitsch case MAT_IGNORE_OFF_PROC_ENTRIES: 929d71ae5a4SJacob Faibussowitsch a->donotstash = flg; 930d71ae5a4SJacob Faibussowitsch break; 93177e54ba9SKris Buschelman case MAT_SYMMETRIC: 93277e54ba9SKris Buschelman case MAT_STRUCTURALLY_SYMMETRIC: 9339a4540c5SBarry Smith case MAT_HERMITIAN: 9349a4540c5SBarry Smith case MAT_SYMMETRY_ETERNAL: 935b94d7dedSBarry Smith case MAT_STRUCTURAL_SYMMETRY_ETERNAL: 936b94d7dedSBarry Smith case MAT_SPD: 937600fe468SBarry Smith case MAT_IGNORE_LOWER_TRIANGULAR: 9385d7aebe8SStefano Zampini case MAT_IGNORE_ZERO_ENTRIES: 939b94d7dedSBarry Smith case MAT_SPD_ETERNAL: 940b94d7dedSBarry Smith /* if the diagonal matrix is square it inherits some of the properties above */ 9419566063dSJacob Faibussowitsch PetscCall(PetscInfo(A, "Option %s ignored\n", MatOptions[op])); 94277e54ba9SKris Buschelman break; 943d71ae5a4SJacob Faibussowitsch default: 944d71ae5a4SJacob Faibussowitsch SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "unknown option %s", MatOptions[op]); 9453a40ed3dSBarry Smith } 9463ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 9478965ea79SLois Curfman McInnes } 9488965ea79SLois Curfman McInnes 94966976f2fSJacob Faibussowitsch static PetscErrorCode MatDiagonalScale_MPIDense(Mat A, Vec ll, Vec rr) 950d71ae5a4SJacob Faibussowitsch { 9515b2fa520SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)A->data; 952637a0070SStefano Zampini const PetscScalar *l; 953637a0070SStefano Zampini PetscScalar x, *v, *vv, *r; 954637a0070SStefano Zampini PetscInt i, j, s2a, s3a, s2, s3, m = mdn->A->rmap->n, n = mdn->A->cmap->n, lda; 9555b2fa520SLois Curfman McInnes 9565b2fa520SLois Curfman McInnes PetscFunctionBegin; 9579566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(mdn->A, &vv)); 9589566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(mdn->A, &lda)); 9599566063dSJacob Faibussowitsch PetscCall(MatGetLocalSize(A, &s2, &s3)); 9605b2fa520SLois Curfman McInnes if (ll) { 9619566063dSJacob Faibussowitsch PetscCall(VecGetLocalSize(ll, &s2a)); 96208401ef6SPierre Jolivet PetscCheck(s2a == s2, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Left scaling vector non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT, s2a, s2); 9639566063dSJacob Faibussowitsch PetscCall(VecGetArrayRead(ll, &l)); 9645b2fa520SLois Curfman McInnes for (i = 0; i < m; i++) { 9655b2fa520SLois Curfman McInnes x = l[i]; 966637a0070SStefano Zampini v = vv + i; 9679371c9d4SSatish Balay for (j = 0; j < n; j++) { 9689371c9d4SSatish Balay (*v) *= x; 9699371c9d4SSatish Balay v += lda; 9709371c9d4SSatish Balay } 9715b2fa520SLois Curfman McInnes } 9729566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayRead(ll, &l)); 9739566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(1.0 * n * m)); 9745b2fa520SLois Curfman McInnes } 9755b2fa520SLois Curfman McInnes if (rr) { 976637a0070SStefano Zampini const PetscScalar *ar; 977637a0070SStefano Zampini 9789566063dSJacob Faibussowitsch PetscCall(VecGetLocalSize(rr, &s3a)); 97908401ef6SPierre Jolivet PetscCheck(s3a == s3, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Right scaling vec non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT ".", s3a, s3); 9809566063dSJacob Faibussowitsch PetscCall(VecGetArrayRead(rr, &ar)); 9816f08ad05SPierre Jolivet if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A)); 9829566063dSJacob Faibussowitsch PetscCall(VecGetArray(mdn->lvec, &r)); 9839566063dSJacob Faibussowitsch PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, ar, r, MPI_REPLACE)); 9849566063dSJacob Faibussowitsch PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ar, r, MPI_REPLACE)); 9859566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayRead(rr, &ar)); 9865b2fa520SLois Curfman McInnes for (i = 0; i < n; i++) { 9875b2fa520SLois Curfman McInnes x = r[i]; 988637a0070SStefano Zampini v = vv + i * lda; 9892205254eSKarl Rupp for (j = 0; j < m; j++) (*v++) *= x; 9905b2fa520SLois Curfman McInnes } 9919566063dSJacob Faibussowitsch PetscCall(VecRestoreArray(mdn->lvec, &r)); 9929566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(1.0 * n * m)); 9935b2fa520SLois Curfman McInnes } 9949566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(mdn->A, &vv)); 9953ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 9965b2fa520SLois Curfman McInnes } 9975b2fa520SLois Curfman McInnes 99866976f2fSJacob Faibussowitsch static PetscErrorCode MatNorm_MPIDense(Mat A, NormType type, PetscReal *nrm) 999d71ae5a4SJacob Faibussowitsch { 10003501a2bdSLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)A->data; 100113f74950SBarry Smith PetscInt i, j; 1002616b8fbbSStefano Zampini PetscMPIInt size; 1003329f5518SBarry Smith PetscReal sum = 0.0; 1004637a0070SStefano Zampini const PetscScalar *av, *v; 10053501a2bdSLois Curfman McInnes 10063a40ed3dSBarry Smith PetscFunctionBegin; 10079566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(mdn->A, &av)); 1008637a0070SStefano Zampini v = av; 10099566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size)); 1010616b8fbbSStefano Zampini if (size == 1) { 10119566063dSJacob Faibussowitsch PetscCall(MatNorm(mdn->A, type, nrm)); 10123501a2bdSLois Curfman McInnes } else { 10133501a2bdSLois Curfman McInnes if (type == NORM_FROBENIUS) { 1014d0f46423SBarry Smith for (i = 0; i < mdn->A->cmap->n * mdn->A->rmap->n; i++) { 10159371c9d4SSatish Balay sum += PetscRealPart(PetscConj(*v) * (*v)); 10169371c9d4SSatish Balay v++; 10173501a2bdSLois Curfman McInnes } 10181c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&sum, nrm, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A))); 10198f1a2a5eSBarry Smith *nrm = PetscSqrtReal(*nrm); 10209566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0 * mdn->A->cmap->n * mdn->A->rmap->n)); 10213a40ed3dSBarry Smith } else if (type == NORM_1) { 1022329f5518SBarry Smith PetscReal *tmp, *tmp2; 10239566063dSJacob Faibussowitsch PetscCall(PetscCalloc2(A->cmap->N, &tmp, A->cmap->N, &tmp2)); 1024064f8208SBarry Smith *nrm = 0.0; 1025637a0070SStefano Zampini v = av; 1026d0f46423SBarry Smith for (j = 0; j < mdn->A->cmap->n; j++) { 1027d0f46423SBarry Smith for (i = 0; i < mdn->A->rmap->n; i++) { 10289371c9d4SSatish Balay tmp[j] += PetscAbsScalar(*v); 10299371c9d4SSatish Balay v++; 10303501a2bdSLois Curfman McInnes } 10313501a2bdSLois Curfman McInnes } 10321c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(tmp, tmp2, A->cmap->N, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A))); 1033d0f46423SBarry Smith for (j = 0; j < A->cmap->N; j++) { 1034064f8208SBarry Smith if (tmp2[j] > *nrm) *nrm = tmp2[j]; 10353501a2bdSLois Curfman McInnes } 10369566063dSJacob Faibussowitsch PetscCall(PetscFree2(tmp, tmp2)); 10379566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(A->cmap->n * A->rmap->n)); 10383a40ed3dSBarry Smith } else if (type == NORM_INFINITY) { /* max row norm */ 1039329f5518SBarry Smith PetscReal ntemp; 10409566063dSJacob Faibussowitsch PetscCall(MatNorm(mdn->A, type, &ntemp)); 10411c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&ntemp, nrm, 1, MPIU_REAL, MPIU_MAX, PetscObjectComm((PetscObject)A))); 1042ce94432eSBarry Smith } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "No support for two norm"); 10433501a2bdSLois Curfman McInnes } 10449566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(mdn->A, &av)); 10453ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 10463501a2bdSLois Curfman McInnes } 10473501a2bdSLois Curfman McInnes 104866976f2fSJacob Faibussowitsch static PetscErrorCode MatTranspose_MPIDense(Mat A, MatReuse reuse, Mat *matout) 1049d71ae5a4SJacob Faibussowitsch { 10503501a2bdSLois Curfman McInnes Mat_MPIDense *a = (Mat_MPIDense *)A->data; 10513501a2bdSLois Curfman McInnes Mat B; 1052d0f46423SBarry Smith PetscInt M = A->rmap->N, N = A->cmap->N, m, n, *rwork, rstart = A->rmap->rstart; 1053637a0070SStefano Zampini PetscInt j, i, lda; 105487828ca2SBarry Smith PetscScalar *v; 10553501a2bdSLois Curfman McInnes 10563a40ed3dSBarry Smith PetscFunctionBegin; 10577fb60732SBarry Smith if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(A, *matout)); 1058cf37664fSBarry Smith if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX) { 10599566063dSJacob Faibussowitsch PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B)); 10609566063dSJacob Faibussowitsch PetscCall(MatSetSizes(B, A->cmap->n, A->rmap->n, N, M)); 10619566063dSJacob Faibussowitsch PetscCall(MatSetType(B, ((PetscObject)A)->type_name)); 10629566063dSJacob Faibussowitsch PetscCall(MatMPIDenseSetPreallocation(B, NULL)); 1063637a0070SStefano Zampini } else B = *matout; 10643501a2bdSLois Curfman McInnes 10659371c9d4SSatish Balay m = a->A->rmap->n; 10669371c9d4SSatish Balay n = a->A->cmap->n; 10679566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)&v)); 10689566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &lda)); 10699566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(m, &rwork)); 10703501a2bdSLois Curfman McInnes for (i = 0; i < m; i++) rwork[i] = rstart + i; 10711acff37aSSatish Balay for (j = 0; j < n; j++) { 10729566063dSJacob Faibussowitsch PetscCall(MatSetValues(B, 1, &j, m, rwork, v, INSERT_VALUES)); 10738e3a54c0SPierre Jolivet v = PetscSafePointerPlusOffset(v, lda); 10743501a2bdSLois Curfman McInnes } 10759566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)&v)); 10769566063dSJacob Faibussowitsch PetscCall(PetscFree(rwork)); 10779566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY)); 10789566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY)); 1079cf37664fSBarry Smith if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_REUSE_MATRIX) { 10803501a2bdSLois Curfman McInnes *matout = B; 10813501a2bdSLois Curfman McInnes } else { 10829566063dSJacob Faibussowitsch PetscCall(MatHeaderMerge(A, &B)); 10833501a2bdSLois Curfman McInnes } 10843ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1085096963f5SLois Curfman McInnes } 1086096963f5SLois Curfman McInnes 10876849ba73SBarry Smith static PetscErrorCode MatDuplicate_MPIDense(Mat, MatDuplicateOption, Mat *); 108852c5f739Sprj- PETSC_INTERN PetscErrorCode MatScale_MPIDense(Mat, PetscScalar); 10898965ea79SLois Curfman McInnes 109066976f2fSJacob Faibussowitsch static PetscErrorCode MatSetUp_MPIDense(Mat A) 1091d71ae5a4SJacob Faibussowitsch { 1092273d9f13SBarry Smith PetscFunctionBegin; 10939566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(A->rmap)); 10949566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(A->cmap)); 109548a46eb9SPierre Jolivet if (!A->preallocated) PetscCall(MatMPIDenseSetPreallocation(A, NULL)); 10963ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1097273d9f13SBarry Smith } 1098273d9f13SBarry Smith 109966976f2fSJacob Faibussowitsch static PetscErrorCode MatAXPY_MPIDense(Mat Y, PetscScalar alpha, Mat X, MatStructure str) 1100d71ae5a4SJacob Faibussowitsch { 1101488007eeSBarry Smith Mat_MPIDense *A = (Mat_MPIDense *)Y->data, *B = (Mat_MPIDense *)X->data; 1102488007eeSBarry Smith 1103488007eeSBarry Smith PetscFunctionBegin; 11049566063dSJacob Faibussowitsch PetscCall(MatAXPY(A->A, alpha, B->A, str)); 11053ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1106488007eeSBarry Smith } 1107488007eeSBarry Smith 110866976f2fSJacob Faibussowitsch static PetscErrorCode MatConjugate_MPIDense(Mat mat) 1109d71ae5a4SJacob Faibussowitsch { 1110ba337c44SJed Brown Mat_MPIDense *a = (Mat_MPIDense *)mat->data; 1111ba337c44SJed Brown 1112ba337c44SJed Brown PetscFunctionBegin; 11139566063dSJacob Faibussowitsch PetscCall(MatConjugate(a->A)); 11143ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1115ba337c44SJed Brown } 1116ba337c44SJed Brown 111766976f2fSJacob Faibussowitsch static PetscErrorCode MatRealPart_MPIDense(Mat A) 1118d71ae5a4SJacob Faibussowitsch { 1119ba337c44SJed Brown Mat_MPIDense *a = (Mat_MPIDense *)A->data; 1120ba337c44SJed Brown 1121ba337c44SJed Brown PetscFunctionBegin; 11229566063dSJacob Faibussowitsch PetscCall(MatRealPart(a->A)); 11233ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1124ba337c44SJed Brown } 1125ba337c44SJed Brown 112666976f2fSJacob Faibussowitsch static PetscErrorCode MatImaginaryPart_MPIDense(Mat A) 1127d71ae5a4SJacob Faibussowitsch { 1128ba337c44SJed Brown Mat_MPIDense *a = (Mat_MPIDense *)A->data; 1129ba337c44SJed Brown 1130ba337c44SJed Brown PetscFunctionBegin; 11319566063dSJacob Faibussowitsch PetscCall(MatImaginaryPart(a->A)); 11323ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1133ba337c44SJed Brown } 1134ba337c44SJed Brown 1135d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatGetColumnVector_MPIDense(Mat A, Vec v, PetscInt col) 1136d71ae5a4SJacob Faibussowitsch { 1137637a0070SStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data; 113849a6ff4bSBarry Smith 113949a6ff4bSBarry Smith PetscFunctionBegin; 114028b400f6SJacob Faibussowitsch PetscCheck(a->A, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing local matrix"); 114128b400f6SJacob Faibussowitsch PetscCheck(a->A->ops->getcolumnvector, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing get column operation"); 11429566063dSJacob Faibussowitsch PetscCall((*a->A->ops->getcolumnvector)(a->A, v, col)); 11433ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 114449a6ff4bSBarry Smith } 114549a6ff4bSBarry Smith 1146857cbf51SRichard Tran Mills PETSC_INTERN PetscErrorCode MatGetColumnReductions_SeqDense(Mat, PetscInt, PetscReal *); 114752c5f739Sprj- 114866976f2fSJacob Faibussowitsch static PetscErrorCode MatGetColumnReductions_MPIDense(Mat A, PetscInt type, PetscReal *reductions) 1149d71ae5a4SJacob Faibussowitsch { 1150a873a8cdSSam Reynolds PetscInt i, m, n; 11510716a85fSBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data; 11520716a85fSBarry Smith PetscReal *work; 11530716a85fSBarry Smith 11540716a85fSBarry Smith PetscFunctionBegin; 11559566063dSJacob Faibussowitsch PetscCall(MatGetSize(A, &m, &n)); 11569566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(n, &work)); 1157857cbf51SRichard Tran Mills if (type == REDUCTION_MEAN_REALPART) { 11589566063dSJacob Faibussowitsch PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_REALPART, work)); 1159857cbf51SRichard Tran Mills } else if (type == REDUCTION_MEAN_IMAGINARYPART) { 11609566063dSJacob Faibussowitsch PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_IMAGINARYPART, work)); 1161a873a8cdSSam Reynolds } else { 11629566063dSJacob Faibussowitsch PetscCall(MatGetColumnReductions_SeqDense(a->A, type, work)); 1163a873a8cdSSam Reynolds } 1164857cbf51SRichard Tran Mills if (type == NORM_2) { 11650716a85fSBarry Smith for (i = 0; i < n; i++) work[i] *= work[i]; 11660716a85fSBarry Smith } 1167857cbf51SRichard Tran Mills if (type == NORM_INFINITY) { 11681c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(work, reductions, n, MPIU_REAL, MPIU_MAX, A->hdr.comm)); 11690716a85fSBarry Smith } else { 11701c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(work, reductions, n, MPIU_REAL, MPIU_SUM, A->hdr.comm)); 11710716a85fSBarry Smith } 11729566063dSJacob Faibussowitsch PetscCall(PetscFree(work)); 1173857cbf51SRichard Tran Mills if (type == NORM_2) { 1174a873a8cdSSam Reynolds for (i = 0; i < n; i++) reductions[i] = PetscSqrtReal(reductions[i]); 1175857cbf51SRichard Tran Mills } else if (type == REDUCTION_MEAN_REALPART || type == REDUCTION_MEAN_IMAGINARYPART) { 1176a873a8cdSSam Reynolds for (i = 0; i < n; i++) reductions[i] /= m; 11770716a85fSBarry Smith } 11783ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 11790716a85fSBarry Smith } 11800716a85fSBarry Smith 1181d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatSetRandom_MPIDense(Mat x, PetscRandom rctx) 1182d71ae5a4SJacob Faibussowitsch { 118373a71a0fSBarry Smith Mat_MPIDense *d = (Mat_MPIDense *)x->data; 118473a71a0fSBarry Smith 118573a71a0fSBarry Smith PetscFunctionBegin; 11869566063dSJacob Faibussowitsch PetscCall(MatSetRandom(d->A, rctx)); 11873faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE) 11883faff063SStefano Zampini x->offloadmask = d->A->offloadmask; 11893faff063SStefano Zampini #endif 11903ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 119173a71a0fSBarry Smith } 119273a71a0fSBarry Smith 1193d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMissingDiagonal_MPIDense(Mat A, PetscBool *missing, PetscInt *d) 1194d71ae5a4SJacob Faibussowitsch { 11953b49f96aSBarry Smith PetscFunctionBegin; 11963b49f96aSBarry Smith *missing = PETSC_FALSE; 11973ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 11983b49f96aSBarry Smith } 11993b49f96aSBarry Smith 12004222ddf1SHong Zhang static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat); 1201cc48ffa7SToby Isaac static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat); 12026718818eSStefano Zampini static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat); 1203a1f9a5e6SStefano Zampini static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat); 12046718818eSStefano Zampini static PetscErrorCode MatEqual_MPIDense(Mat, Mat, PetscBool *); 12056718818eSStefano Zampini static PetscErrorCode MatLoad_MPIDense(Mat, PetscViewer); 12065d8c7819SPierre Jolivet static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat); 1207cc48ffa7SToby Isaac 120809dc0095SBarry Smith static struct _MatOps MatOps_Values = {MatSetValues_MPIDense, 120909dc0095SBarry Smith MatGetRow_MPIDense, 121009dc0095SBarry Smith MatRestoreRow_MPIDense, 121109dc0095SBarry Smith MatMult_MPIDense, 121297304618SKris Buschelman /* 4*/ MatMultAdd_MPIDense, 12137c922b88SBarry Smith MatMultTranspose_MPIDense, 12147c922b88SBarry Smith MatMultTransposeAdd_MPIDense, 1215f4259b30SLisandro Dalcin NULL, 1216f4259b30SLisandro Dalcin NULL, 1217f4259b30SLisandro Dalcin NULL, 1218f4259b30SLisandro Dalcin /* 10*/ NULL, 1219f4259b30SLisandro Dalcin NULL, 1220f4259b30SLisandro Dalcin NULL, 1221f4259b30SLisandro Dalcin NULL, 122209dc0095SBarry Smith MatTranspose_MPIDense, 122397304618SKris Buschelman /* 15*/ MatGetInfo_MPIDense, 12246e4ee0c6SHong Zhang MatEqual_MPIDense, 122509dc0095SBarry Smith MatGetDiagonal_MPIDense, 12265b2fa520SLois Curfman McInnes MatDiagonalScale_MPIDense, 122709dc0095SBarry Smith MatNorm_MPIDense, 122897304618SKris Buschelman /* 20*/ MatAssemblyBegin_MPIDense, 122909dc0095SBarry Smith MatAssemblyEnd_MPIDense, 123009dc0095SBarry Smith MatSetOption_MPIDense, 123109dc0095SBarry Smith MatZeroEntries_MPIDense, 1232d519adbfSMatthew Knepley /* 24*/ MatZeroRows_MPIDense, 1233f4259b30SLisandro Dalcin NULL, 1234f4259b30SLisandro Dalcin NULL, 1235f4259b30SLisandro Dalcin NULL, 1236f4259b30SLisandro Dalcin NULL, 12374994cf47SJed Brown /* 29*/ MatSetUp_MPIDense, 1238f4259b30SLisandro Dalcin NULL, 1239f4259b30SLisandro Dalcin NULL, 1240c56a70eeSBarry Smith MatGetDiagonalBlock_MPIDense, 1241f4259b30SLisandro Dalcin NULL, 1242d519adbfSMatthew Knepley /* 34*/ MatDuplicate_MPIDense, 1243f4259b30SLisandro Dalcin NULL, 1244f4259b30SLisandro Dalcin NULL, 1245f4259b30SLisandro Dalcin NULL, 1246f4259b30SLisandro Dalcin NULL, 1247d519adbfSMatthew Knepley /* 39*/ MatAXPY_MPIDense, 12487dae84e0SHong Zhang MatCreateSubMatrices_MPIDense, 1249f4259b30SLisandro Dalcin NULL, 125009dc0095SBarry Smith MatGetValues_MPIDense, 1251a76f77c3SStefano Zampini MatCopy_MPIDense, 1252f4259b30SLisandro Dalcin /* 44*/ NULL, 125309dc0095SBarry Smith MatScale_MPIDense, 12542f605a99SJose E. Roman MatShift_MPIDense, 1255f4259b30SLisandro Dalcin NULL, 1256f4259b30SLisandro Dalcin NULL, 125773a71a0fSBarry Smith /* 49*/ MatSetRandom_MPIDense, 1258f4259b30SLisandro Dalcin NULL, 1259f4259b30SLisandro Dalcin NULL, 1260f4259b30SLisandro Dalcin NULL, 1261f4259b30SLisandro Dalcin NULL, 1262f4259b30SLisandro Dalcin /* 54*/ NULL, 1263f4259b30SLisandro Dalcin NULL, 1264f4259b30SLisandro Dalcin NULL, 1265f4259b30SLisandro Dalcin NULL, 1266f4259b30SLisandro Dalcin NULL, 12677dae84e0SHong Zhang /* 59*/ MatCreateSubMatrix_MPIDense, 1268b9b97703SBarry Smith MatDestroy_MPIDense, 1269b9b97703SBarry Smith MatView_MPIDense, 1270f4259b30SLisandro Dalcin NULL, 1271f4259b30SLisandro Dalcin NULL, 1272f4259b30SLisandro Dalcin /* 64*/ NULL, 1273f4259b30SLisandro Dalcin NULL, 1274f4259b30SLisandro Dalcin NULL, 1275f4259b30SLisandro Dalcin NULL, 1276f4259b30SLisandro Dalcin NULL, 1277f4259b30SLisandro Dalcin /* 69*/ NULL, 1278f4259b30SLisandro Dalcin NULL, 1279f4259b30SLisandro Dalcin NULL, 1280f4259b30SLisandro Dalcin NULL, 1281f4259b30SLisandro Dalcin NULL, 1282f4259b30SLisandro Dalcin /* 74*/ NULL, 1283f4259b30SLisandro Dalcin NULL, 1284f4259b30SLisandro Dalcin NULL, 1285f4259b30SLisandro Dalcin NULL, 1286f4259b30SLisandro Dalcin NULL, 1287f4259b30SLisandro Dalcin /* 79*/ NULL, 1288f4259b30SLisandro Dalcin NULL, 1289f4259b30SLisandro Dalcin NULL, 1290f4259b30SLisandro Dalcin NULL, 12915bba2384SShri Abhyankar /* 83*/ MatLoad_MPIDense, 1292f4259b30SLisandro Dalcin NULL, 1293f4259b30SLisandro Dalcin NULL, 1294f4259b30SLisandro Dalcin NULL, 1295f4259b30SLisandro Dalcin NULL, 1296f4259b30SLisandro Dalcin NULL, 1297f4259b30SLisandro Dalcin /* 89*/ NULL, 1298f4259b30SLisandro Dalcin NULL, 1299f4259b30SLisandro Dalcin NULL, 1300f4259b30SLisandro Dalcin NULL, 1301f4259b30SLisandro Dalcin NULL, 1302f4259b30SLisandro Dalcin /* 94*/ NULL, 1303f4259b30SLisandro Dalcin NULL, 13046718818eSStefano Zampini MatMatTransposeMultSymbolic_MPIDense_MPIDense, 1305cc48ffa7SToby Isaac MatMatTransposeMultNumeric_MPIDense_MPIDense, 1306f4259b30SLisandro Dalcin NULL, 13074222ddf1SHong Zhang /* 99*/ MatProductSetFromOptions_MPIDense, 1308f4259b30SLisandro Dalcin NULL, 1309f4259b30SLisandro Dalcin NULL, 1310ba337c44SJed Brown MatConjugate_MPIDense, 1311f4259b30SLisandro Dalcin NULL, 1312f4259b30SLisandro Dalcin /*104*/ NULL, 1313ba337c44SJed Brown MatRealPart_MPIDense, 1314ba337c44SJed Brown MatImaginaryPart_MPIDense, 1315f4259b30SLisandro Dalcin NULL, 1316f4259b30SLisandro Dalcin NULL, 1317f4259b30SLisandro Dalcin /*109*/ NULL, 1318f4259b30SLisandro Dalcin NULL, 1319f4259b30SLisandro Dalcin NULL, 132049a6ff4bSBarry Smith MatGetColumnVector_MPIDense, 13213b49f96aSBarry Smith MatMissingDiagonal_MPIDense, 1322f4259b30SLisandro Dalcin /*114*/ NULL, 1323f4259b30SLisandro Dalcin NULL, 1324f4259b30SLisandro Dalcin NULL, 1325f4259b30SLisandro Dalcin NULL, 1326f4259b30SLisandro Dalcin NULL, 1327f4259b30SLisandro Dalcin /*119*/ NULL, 1328f4259b30SLisandro Dalcin NULL, 1329459e8d23SBlanca Mellado Pinto MatMultHermitianTranspose_MPIDense, 1330459e8d23SBlanca Mellado Pinto MatMultHermitianTransposeAdd_MPIDense, 1331f4259b30SLisandro Dalcin NULL, 1332f4259b30SLisandro Dalcin /*124*/ NULL, 1333a873a8cdSSam Reynolds MatGetColumnReductions_MPIDense, 1334f4259b30SLisandro Dalcin NULL, 1335f4259b30SLisandro Dalcin NULL, 1336f4259b30SLisandro Dalcin NULL, 1337f4259b30SLisandro Dalcin /*129*/ NULL, 1338f4259b30SLisandro Dalcin NULL, 13396718818eSStefano Zampini MatTransposeMatMultSymbolic_MPIDense_MPIDense, 1340cb20be35SHong Zhang MatTransposeMatMultNumeric_MPIDense_MPIDense, 1341f4259b30SLisandro Dalcin NULL, 1342f4259b30SLisandro Dalcin /*134*/ NULL, 1343f4259b30SLisandro Dalcin NULL, 1344f4259b30SLisandro Dalcin NULL, 1345f4259b30SLisandro Dalcin NULL, 1346f4259b30SLisandro Dalcin NULL, 1347f4259b30SLisandro Dalcin /*139*/ NULL, 1348f4259b30SLisandro Dalcin NULL, 1349f4259b30SLisandro Dalcin NULL, 1350f4259b30SLisandro Dalcin NULL, 1351f4259b30SLisandro Dalcin NULL, 13524222ddf1SHong Zhang MatCreateMPIMatConcatenateSeqMat_MPIDense, 1353f4259b30SLisandro Dalcin /*145*/ NULL, 1354f4259b30SLisandro Dalcin NULL, 135599a7f59eSMark Adams NULL, 135699a7f59eSMark Adams NULL, 13577fb60732SBarry Smith NULL, 1358dec0b466SHong Zhang /*150*/ NULL, 1359eede4a3fSMark Adams NULL, 13604cc2b5b5SPierre Jolivet NULL, 1361*42ce410bSJunchao Zhang NULL, 1362*42ce410bSJunchao Zhang NULL, 1363dec0b466SHong Zhang NULL}; 13648965ea79SLois Curfman McInnes 136566976f2fSJacob Faibussowitsch static PetscErrorCode MatMPIDenseSetPreallocation_MPIDense(Mat mat, PetscScalar *data) 1366d71ae5a4SJacob Faibussowitsch { 1367637a0070SStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)mat->data; 136847d993e7Ssuyashtn MatType mtype = MATSEQDENSE; 1369a23d5eceSKris Buschelman 1370a23d5eceSKris Buschelman PetscFunctionBegin; 137128b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 13729566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(mat->rmap)); 13739566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(mat->cmap)); 1374637a0070SStefano Zampini if (!a->A) { 13759566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF, &a->A)); 13769566063dSJacob Faibussowitsch PetscCall(MatSetSizes(a->A, mat->rmap->n, mat->cmap->N, mat->rmap->n, mat->cmap->N)); 1377637a0070SStefano Zampini } 137847d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) 137947d993e7Ssuyashtn PetscBool iscuda; 13809566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSECUDA, &iscuda)); 138147d993e7Ssuyashtn if (iscuda) mtype = MATSEQDENSECUDA; 138247d993e7Ssuyashtn #endif 138347d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP) 138447d993e7Ssuyashtn PetscBool iship; 138547d993e7Ssuyashtn PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSEHIP, &iship)); 138647d993e7Ssuyashtn if (iship) mtype = MATSEQDENSEHIP; 138747d993e7Ssuyashtn #endif 138847d993e7Ssuyashtn PetscCall(MatSetType(a->A, mtype)); 13899566063dSJacob Faibussowitsch PetscCall(MatSeqDenseSetPreallocation(a->A, data)); 139047d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 13913faff063SStefano Zampini mat->offloadmask = a->A->offloadmask; 13923faff063SStefano Zampini #endif 1393637a0070SStefano Zampini mat->preallocated = PETSC_TRUE; 13946f08ad05SPierre Jolivet mat->assembled = PETSC_TRUE; 13953ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1396a23d5eceSKris Buschelman } 1397a23d5eceSKris Buschelman 1398d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode MatConvert_MPIAIJ_MPIDense(Mat A, MatType newtype, MatReuse reuse, Mat *newmat) 1399d71ae5a4SJacob Faibussowitsch { 14004e29119aSPierre Jolivet Mat B, C; 14014e29119aSPierre Jolivet 14024e29119aSPierre Jolivet PetscFunctionBegin; 14039566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A, MAT_INITIAL_MATRIX, &C)); 14049566063dSJacob Faibussowitsch PetscCall(MatConvert_SeqAIJ_SeqDense(C, MATSEQDENSE, MAT_INITIAL_MATRIX, &B)); 14059566063dSJacob Faibussowitsch PetscCall(MatDestroy(&C)); 14064e29119aSPierre Jolivet if (reuse == MAT_REUSE_MATRIX) { 14074e29119aSPierre Jolivet C = *newmat; 14084e29119aSPierre Jolivet } else C = NULL; 14099566063dSJacob Faibussowitsch PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C)); 14109566063dSJacob Faibussowitsch PetscCall(MatDestroy(&B)); 14114e29119aSPierre Jolivet if (reuse == MAT_INPLACE_MATRIX) { 14129566063dSJacob Faibussowitsch PetscCall(MatHeaderReplace(A, &C)); 14134e29119aSPierre Jolivet } else if (reuse == MAT_INITIAL_MATRIX) *newmat = C; 14143ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 14154e29119aSPierre Jolivet } 14164e29119aSPierre Jolivet 141766976f2fSJacob Faibussowitsch static PetscErrorCode MatConvert_MPIDense_MPIAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat) 1418d71ae5a4SJacob Faibussowitsch { 14194e29119aSPierre Jolivet Mat B, C; 14204e29119aSPierre Jolivet 14214e29119aSPierre Jolivet PetscFunctionBegin; 14229566063dSJacob Faibussowitsch PetscCall(MatDenseGetLocalMatrix(A, &C)); 14239566063dSJacob Faibussowitsch PetscCall(MatConvert_SeqDense_SeqAIJ(C, MATSEQAIJ, MAT_INITIAL_MATRIX, &B)); 14244e29119aSPierre Jolivet if (reuse == MAT_REUSE_MATRIX) { 14254e29119aSPierre Jolivet C = *newmat; 14264e29119aSPierre Jolivet } else C = NULL; 14279566063dSJacob Faibussowitsch PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C)); 14289566063dSJacob Faibussowitsch PetscCall(MatDestroy(&B)); 14294e29119aSPierre Jolivet if (reuse == MAT_INPLACE_MATRIX) { 14309566063dSJacob Faibussowitsch PetscCall(MatHeaderReplace(A, &C)); 14314e29119aSPierre Jolivet } else if (reuse == MAT_INITIAL_MATRIX) *newmat = C; 14323ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 14334e29119aSPierre Jolivet } 14344e29119aSPierre Jolivet 143565b80a83SHong Zhang #if defined(PETSC_HAVE_ELEMENTAL) 1436d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode MatConvert_MPIDense_Elemental(Mat A, MatType newtype, MatReuse reuse, Mat *newmat) 1437d71ae5a4SJacob Faibussowitsch { 1438f709e8d7SPierre Jolivet Mat_MPIDense *a = (Mat_MPIDense *)A->data; 14398ea901baSHong Zhang Mat mat_elemental; 144032d7a744SHong Zhang PetscScalar *v; 1441f709e8d7SPierre Jolivet PetscInt m = A->rmap->n, N = A->cmap->N, rstart = A->rmap->rstart, i, *rows, *cols, lda; 14428ea901baSHong Zhang 14438baccfbdSHong Zhang PetscFunctionBegin; 1444378336b6SHong Zhang if (reuse == MAT_REUSE_MATRIX) { 1445378336b6SHong Zhang mat_elemental = *newmat; 14469566063dSJacob Faibussowitsch PetscCall(MatZeroEntries(*newmat)); 1447378336b6SHong Zhang } else { 14489566063dSJacob Faibussowitsch PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat_elemental)); 14499566063dSJacob Faibussowitsch PetscCall(MatSetSizes(mat_elemental, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N)); 14509566063dSJacob Faibussowitsch PetscCall(MatSetType(mat_elemental, MATELEMENTAL)); 14519566063dSJacob Faibussowitsch PetscCall(MatSetUp(mat_elemental)); 14529566063dSJacob Faibussowitsch PetscCall(MatSetOption(mat_elemental, MAT_ROW_ORIENTED, PETSC_FALSE)); 1453378336b6SHong Zhang } 1454378336b6SHong Zhang 14559566063dSJacob Faibussowitsch PetscCall(PetscMalloc2(m, &rows, N, &cols)); 145632d7a744SHong Zhang for (i = 0; i < N; i++) cols[i] = i; 145732d7a744SHong Zhang for (i = 0; i < m; i++) rows[i] = rstart + i; 14588ea901baSHong Zhang 1459637a0070SStefano Zampini /* PETSc-Elemental interface uses axpy for setting off-processor entries, only ADD_VALUES is allowed */ 14609566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(A, &v)); 1461f709e8d7SPierre Jolivet PetscCall(MatDenseGetLDA(a->A, &lda)); 1462f709e8d7SPierre Jolivet if (lda == m) PetscCall(MatSetValues(mat_elemental, m, rows, N, cols, v, ADD_VALUES)); 1463f709e8d7SPierre Jolivet else { 1464f709e8d7SPierre Jolivet for (i = 0; i < N; i++) PetscCall(MatSetValues(mat_elemental, m, rows, 1, &i, v + lda * i, ADD_VALUES)); 1465f709e8d7SPierre Jolivet } 14669566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(mat_elemental, MAT_FINAL_ASSEMBLY)); 14679566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(mat_elemental, MAT_FINAL_ASSEMBLY)); 14689566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(A, &v)); 14699566063dSJacob Faibussowitsch PetscCall(PetscFree2(rows, cols)); 14708ea901baSHong Zhang 1471511c6705SHong Zhang if (reuse == MAT_INPLACE_MATRIX) { 14729566063dSJacob Faibussowitsch PetscCall(MatHeaderReplace(A, &mat_elemental)); 14738ea901baSHong Zhang } else { 14748ea901baSHong Zhang *newmat = mat_elemental; 14758ea901baSHong Zhang } 14763ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 14778baccfbdSHong Zhang } 147865b80a83SHong Zhang #endif 14798baccfbdSHong Zhang 1480d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetColumn_MPIDense(Mat A, PetscInt col, PetscScalar **vals) 1481d71ae5a4SJacob Faibussowitsch { 148286aefd0dSHong Zhang Mat_MPIDense *mat = (Mat_MPIDense *)A->data; 148386aefd0dSHong Zhang 148486aefd0dSHong Zhang PetscFunctionBegin; 14859566063dSJacob Faibussowitsch PetscCall(MatDenseGetColumn(mat->A, col, vals)); 14863ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 148786aefd0dSHong Zhang } 148886aefd0dSHong Zhang 1489d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseRestoreColumn_MPIDense(Mat A, PetscScalar **vals) 1490d71ae5a4SJacob Faibussowitsch { 149186aefd0dSHong Zhang Mat_MPIDense *mat = (Mat_MPIDense *)A->data; 149286aefd0dSHong Zhang 149386aefd0dSHong Zhang PetscFunctionBegin; 14949566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreColumn(mat->A, vals)); 14953ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 149686aefd0dSHong Zhang } 149786aefd0dSHong Zhang 1498d71ae5a4SJacob Faibussowitsch PetscErrorCode MatCreateMPIMatConcatenateSeqMat_MPIDense(MPI_Comm comm, Mat inmat, PetscInt n, MatReuse scall, Mat *outmat) 1499d71ae5a4SJacob Faibussowitsch { 150094e2cb23SJakub Kruzik Mat_MPIDense *mat; 150194e2cb23SJakub Kruzik PetscInt m, nloc, N; 150294e2cb23SJakub Kruzik 150394e2cb23SJakub Kruzik PetscFunctionBegin; 15049566063dSJacob Faibussowitsch PetscCall(MatGetSize(inmat, &m, &N)); 15059566063dSJacob Faibussowitsch PetscCall(MatGetLocalSize(inmat, NULL, &nloc)); 150694e2cb23SJakub Kruzik if (scall == MAT_INITIAL_MATRIX) { /* symbolic phase */ 150794e2cb23SJakub Kruzik PetscInt sum; 150894e2cb23SJakub Kruzik 150948a46eb9SPierre Jolivet if (n == PETSC_DECIDE) PetscCall(PetscSplitOwnership(comm, &n, &N)); 151094e2cb23SJakub Kruzik /* Check sum(n) = N */ 15111c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&n, &sum, 1, MPIU_INT, MPI_SUM, comm)); 151208401ef6SPierre Jolivet PetscCheck(sum == N, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Sum of local columns %" PetscInt_FMT " != global columns %" PetscInt_FMT, sum, N); 151394e2cb23SJakub Kruzik 15149566063dSJacob Faibussowitsch PetscCall(MatCreateDense(comm, m, n, PETSC_DETERMINE, N, NULL, outmat)); 15159566063dSJacob Faibussowitsch PetscCall(MatSetOption(*outmat, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE)); 151694e2cb23SJakub Kruzik } 151794e2cb23SJakub Kruzik 151894e2cb23SJakub Kruzik /* numeric phase */ 151994e2cb23SJakub Kruzik mat = (Mat_MPIDense *)(*outmat)->data; 15209566063dSJacob Faibussowitsch PetscCall(MatCopy(inmat, mat->A, SAME_NONZERO_PATTERN)); 15213ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 152294e2cb23SJakub Kruzik } 152394e2cb23SJakub Kruzik 1524d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVec_MPIDense(Mat A, PetscInt col, Vec *v) 1525d71ae5a4SJacob Faibussowitsch { 15266947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data; 15276947451fSStefano Zampini PetscInt lda; 15286947451fSStefano Zampini 15296947451fSStefano Zampini PetscFunctionBegin; 153028b400f6SJacob Faibussowitsch PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 153128b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 1532d16ceb75SStefano Zampini if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec)); 15336947451fSStefano Zampini a->vecinuse = col + 1; 15349566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &lda)); 15359566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(a->A, (PetscScalar **)&a->ptrinuse)); 15368e3a54c0SPierre Jolivet PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda))); 15376947451fSStefano Zampini *v = a->cvec; 15383ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 15396947451fSStefano Zampini } 15406947451fSStefano Zampini 1541d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVec_MPIDense(Mat A, PetscInt col, Vec *v) 1542d71ae5a4SJacob Faibussowitsch { 15436947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data; 15446947451fSStefano Zampini 15456947451fSStefano Zampini PetscFunctionBegin; 154628b400f6SJacob Faibussowitsch PetscCheck(a->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first"); 154728b400f6SJacob Faibussowitsch PetscCheck(a->cvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column vector"); 15484186a4bbSPierre Jolivet VecCheckAssembled(a->cvec); 15496947451fSStefano Zampini a->vecinuse = 0; 15509566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(a->A, (PetscScalar **)&a->ptrinuse)); 15519566063dSJacob Faibussowitsch PetscCall(VecResetArray(a->cvec)); 1552742765d3SMatthew Knepley if (v) *v = NULL; 15533ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 15546947451fSStefano Zampini } 15556947451fSStefano Zampini 1556d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v) 1557d71ae5a4SJacob Faibussowitsch { 15586947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data; 15596947451fSStefano Zampini PetscInt lda; 15606947451fSStefano Zampini 15616947451fSStefano Zampini PetscFunctionBegin; 156228b400f6SJacob Faibussowitsch PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 156328b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 1564d16ceb75SStefano Zampini if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec)); 15656947451fSStefano Zampini a->vecinuse = col + 1; 15669566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &lda)); 15679566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(a->A, &a->ptrinuse)); 15688e3a54c0SPierre Jolivet PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda))); 15699566063dSJacob Faibussowitsch PetscCall(VecLockReadPush(a->cvec)); 15706947451fSStefano Zampini *v = a->cvec; 15713ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 15726947451fSStefano Zampini } 15736947451fSStefano Zampini 1574d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v) 1575d71ae5a4SJacob Faibussowitsch { 15766947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data; 15776947451fSStefano Zampini 15786947451fSStefano Zampini PetscFunctionBegin; 15795f80ce2aSJacob Faibussowitsch PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first"); 15805f80ce2aSJacob Faibussowitsch PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector"); 15814186a4bbSPierre Jolivet VecCheckAssembled(a->cvec); 15826947451fSStefano Zampini a->vecinuse = 0; 15839566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(a->A, &a->ptrinuse)); 15849566063dSJacob Faibussowitsch PetscCall(VecLockReadPop(a->cvec)); 15859566063dSJacob Faibussowitsch PetscCall(VecResetArray(a->cvec)); 1586742765d3SMatthew Knepley if (v) *v = NULL; 15873ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 15886947451fSStefano Zampini } 15896947451fSStefano Zampini 1590d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v) 1591d71ae5a4SJacob Faibussowitsch { 15926947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data; 15936947451fSStefano Zampini PetscInt lda; 15946947451fSStefano Zampini 15956947451fSStefano Zampini PetscFunctionBegin; 15965f80ce2aSJacob Faibussowitsch PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 15975f80ce2aSJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 1598d16ceb75SStefano Zampini if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec)); 15996947451fSStefano Zampini a->vecinuse = col + 1; 16009566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &lda)); 16019566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayWrite(a->A, (PetscScalar **)&a->ptrinuse)); 16028e3a54c0SPierre Jolivet PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda))); 16036947451fSStefano Zampini *v = a->cvec; 16043ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 16056947451fSStefano Zampini } 16066947451fSStefano Zampini 1607d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v) 1608d71ae5a4SJacob Faibussowitsch { 16096947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data; 16106947451fSStefano Zampini 16116947451fSStefano Zampini PetscFunctionBegin; 16125f80ce2aSJacob Faibussowitsch PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first"); 16135f80ce2aSJacob Faibussowitsch PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector"); 16144186a4bbSPierre Jolivet VecCheckAssembled(a->cvec); 16156947451fSStefano Zampini a->vecinuse = 0; 16169566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayWrite(a->A, (PetscScalar **)&a->ptrinuse)); 16179566063dSJacob Faibussowitsch PetscCall(VecResetArray(a->cvec)); 1618742765d3SMatthew Knepley if (v) *v = NULL; 16193ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 16206947451fSStefano Zampini } 16216947451fSStefano Zampini 162266976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseGetSubMatrix_MPIDense(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *v) 1623d71ae5a4SJacob Faibussowitsch { 16245ea7661aSPierre Jolivet Mat_MPIDense *a = (Mat_MPIDense *)A->data; 1625616b8fbbSStefano Zampini Mat_MPIDense *c; 1626616b8fbbSStefano Zampini MPI_Comm comm; 1627a2748737SPierre Jolivet PetscInt pbegin, pend; 16285ea7661aSPierre Jolivet 16295ea7661aSPierre Jolivet PetscFunctionBegin; 16309566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A, &comm)); 16315f80ce2aSJacob Faibussowitsch PetscCheck(!a->vecinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first"); 16325f80ce2aSJacob Faibussowitsch PetscCheck(!a->matinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 1633a2748737SPierre Jolivet pbegin = PetscMax(0, PetscMin(A->rmap->rend, rbegin) - A->rmap->rstart); 1634a2748737SPierre Jolivet pend = PetscMin(A->rmap->n, PetscMax(0, rend - A->rmap->rstart)); 16355ea7661aSPierre Jolivet if (!a->cmat) { 16369566063dSJacob Faibussowitsch PetscCall(MatCreate(comm, &a->cmat)); 16379566063dSJacob Faibussowitsch PetscCall(MatSetType(a->cmat, ((PetscObject)A)->type_name)); 1638a2748737SPierre Jolivet if (rend - rbegin == A->rmap->N) PetscCall(PetscLayoutReference(A->rmap, &a->cmat->rmap)); 1639a2748737SPierre Jolivet else { 1640a2748737SPierre Jolivet PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, pend - pbegin)); 1641a2748737SPierre Jolivet PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin)); 1642a2748737SPierre Jolivet PetscCall(PetscLayoutSetUp(a->cmat->rmap)); 1643a2748737SPierre Jolivet } 16449566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin)); 16459566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(a->cmat->cmap)); 1646a2748737SPierre Jolivet } else { 1647a2748737SPierre Jolivet PetscBool same = (PetscBool)(rend - rbegin == a->cmat->rmap->N); 1648a2748737SPierre Jolivet if (same && a->cmat->rmap->N != A->rmap->N) { 1649a2748737SPierre Jolivet same = (PetscBool)(pend - pbegin == a->cmat->rmap->n); 1650a2748737SPierre Jolivet PetscCall(MPIU_Allreduce(MPI_IN_PLACE, &same, 1, MPIU_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A))); 1651a2748737SPierre Jolivet } 1652a2748737SPierre Jolivet if (!same) { 1653a2748737SPierre Jolivet PetscCall(PetscLayoutDestroy(&a->cmat->rmap)); 1654a2748737SPierre Jolivet PetscCall(PetscLayoutCreate(comm, &a->cmat->rmap)); 1655a2748737SPierre Jolivet PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, pend - pbegin)); 1656a2748737SPierre Jolivet PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin)); 1657a2748737SPierre Jolivet PetscCall(PetscLayoutSetUp(a->cmat->rmap)); 1658a2748737SPierre Jolivet } 1659a2748737SPierre Jolivet if (cend - cbegin != a->cmat->cmap->N) { 16609566063dSJacob Faibussowitsch PetscCall(PetscLayoutDestroy(&a->cmat->cmap)); 16619566063dSJacob Faibussowitsch PetscCall(PetscLayoutCreate(comm, &a->cmat->cmap)); 16629566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin)); 16639566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(a->cmat->cmap)); 16645ea7661aSPierre Jolivet } 1665a2748737SPierre Jolivet } 1666616b8fbbSStefano Zampini c = (Mat_MPIDense *)a->cmat->data; 16675f80ce2aSJacob Faibussowitsch PetscCheck(!c->A, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first"); 1668a2748737SPierre Jolivet PetscCall(MatDenseGetSubMatrix(a->A, pbegin, pend, cbegin, cend, &c->A)); 16693faff063SStefano Zampini 1670616b8fbbSStefano Zampini a->cmat->preallocated = PETSC_TRUE; 1671616b8fbbSStefano Zampini a->cmat->assembled = PETSC_TRUE; 16723faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE) 16733faff063SStefano Zampini a->cmat->offloadmask = c->A->offloadmask; 16743faff063SStefano Zampini #endif 16755ea7661aSPierre Jolivet a->matinuse = cbegin + 1; 16765ea7661aSPierre Jolivet *v = a->cmat; 16773ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 16785ea7661aSPierre Jolivet } 16795ea7661aSPierre Jolivet 168066976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreSubMatrix_MPIDense(Mat A, Mat *v) 1681d71ae5a4SJacob Faibussowitsch { 16825ea7661aSPierre Jolivet Mat_MPIDense *a = (Mat_MPIDense *)A->data; 1683616b8fbbSStefano Zampini Mat_MPIDense *c; 16845ea7661aSPierre Jolivet 16855ea7661aSPierre Jolivet PetscFunctionBegin; 16865f80ce2aSJacob Faibussowitsch PetscCheck(a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetSubMatrix() first"); 16875f80ce2aSJacob Faibussowitsch PetscCheck(a->cmat, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal matrix"); 16885f80ce2aSJacob Faibussowitsch PetscCheck(*v == a->cmat, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Not the matrix obtained from MatDenseGetSubMatrix()"); 16895ea7661aSPierre Jolivet a->matinuse = 0; 1690616b8fbbSStefano Zampini c = (Mat_MPIDense *)a->cmat->data; 16919566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreSubMatrix(a->A, &c->A)); 1692742765d3SMatthew Knepley if (v) *v = NULL; 16933faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE) 16943faff063SStefano Zampini A->offloadmask = a->A->offloadmask; 16953faff063SStefano Zampini #endif 16963ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 16975ea7661aSPierre Jolivet } 16985ea7661aSPierre Jolivet 1699002bc6daSRichard Tran Mills /*MC 1700002bc6daSRichard Tran Mills MATMPIDENSE - MATMPIDENSE = "mpidense" - A matrix type to be used for distributed dense matrices. 1701002bc6daSRichard Tran Mills 17022ef1f0ffSBarry Smith Options Database Key: 170311a5261eSBarry Smith . -mat_type mpidense - sets the matrix type to `MATMPIDENSE` during a call to `MatSetFromOptions()` 1704002bc6daSRichard Tran Mills 1705002bc6daSRichard Tran Mills Level: beginner 1706002bc6daSRichard Tran Mills 17071cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MatCreateDense()`, `MATSEQDENSE`, `MATDENSE` 1708002bc6daSRichard Tran Mills M*/ 17094742e46bSJacob Faibussowitsch PetscErrorCode MatCreate_MPIDense(Mat mat) 1710d71ae5a4SJacob Faibussowitsch { 1711273d9f13SBarry Smith Mat_MPIDense *a; 1712273d9f13SBarry Smith 1713273d9f13SBarry Smith PetscFunctionBegin; 17144dfa11a4SJacob Faibussowitsch PetscCall(PetscNew(&a)); 1715b0a32e0cSBarry Smith mat->data = (void *)a; 1716aea10558SJacob Faibussowitsch mat->ops[0] = MatOps_Values; 1717273d9f13SBarry Smith 1718273d9f13SBarry Smith mat->insertmode = NOT_SET_VALUES; 1719273d9f13SBarry Smith 1720273d9f13SBarry Smith /* build cache for off array entries formed */ 1721273d9f13SBarry Smith a->donotstash = PETSC_FALSE; 17222205254eSKarl Rupp 17239566063dSJacob Faibussowitsch PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)mat), 1, &mat->stash)); 1724273d9f13SBarry Smith 1725273d9f13SBarry Smith /* stuff used for matrix vector multiply */ 1726f4259b30SLisandro Dalcin a->lvec = NULL; 1727f4259b30SLisandro Dalcin a->Mvctx = NULL; 1728273d9f13SBarry Smith a->roworiented = PETSC_TRUE; 1729273d9f13SBarry Smith 17309566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", MatDenseGetLDA_MPIDense)); 17319566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", MatDenseSetLDA_MPIDense)); 17329566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", MatDenseGetArray_MPIDense)); 17339566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", MatDenseRestoreArray_MPIDense)); 17349566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", MatDenseGetArrayRead_MPIDense)); 17359566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", MatDenseRestoreArrayRead_MPIDense)); 17369566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", MatDenseGetArrayWrite_MPIDense)); 17379566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", MatDenseRestoreArrayWrite_MPIDense)); 17389566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", MatDensePlaceArray_MPIDense)); 17399566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", MatDenseResetArray_MPIDense)); 17409566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", MatDenseReplaceArray_MPIDense)); 17419566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_MPIDense)); 17429566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_MPIDense)); 17439566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_MPIDense)); 17449566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_MPIDense)); 17459566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_MPIDense)); 17469566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_MPIDense)); 17479566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_MPIDense)); 17489566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_MPIDense)); 17499566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", MatConvert_MPIAIJ_MPIDense)); 17509566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", MatConvert_MPIDense_MPIAIJ)); 17518baccfbdSHong Zhang #if defined(PETSC_HAVE_ELEMENTAL) 17529566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", MatConvert_MPIDense_Elemental)); 17538baccfbdSHong Zhang #endif 1754d24d4204SJose E. Roman #if defined(PETSC_HAVE_SCALAPACK) 17559566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", MatConvert_Dense_ScaLAPACK)); 1756d24d4204SJose E. Roman #endif 1757637a0070SStefano Zampini #if defined(PETSC_HAVE_CUDA) 17589566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", MatConvert_MPIDense_MPIDenseCUDA)); 1759637a0070SStefano Zampini #endif 17609566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", MatMPIDenseSetPreallocation_MPIDense)); 17619566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense)); 17629566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", MatProductSetFromOptions_MPIDense_MPIAIJ)); 17636718818eSStefano Zampini #if defined(PETSC_HAVE_CUDA) 17649566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense)); 17659566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ)); 17666718818eSStefano Zampini #endif 176747d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP) 176847d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", MatConvert_MPIDense_MPIDenseHIP)); 176947d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense)); 177047d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ)); 177147d993e7Ssuyashtn #endif 17729566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", MatDenseGetColumn_MPIDense)); 17739566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", MatDenseRestoreColumn_MPIDense)); 17740be0d8bdSHansol Suh PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", MatMultAddColumnRange_MPIDense)); 17750be0d8bdSHansol Suh PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", MatMultHermitianTransposeColumnRange_MPIDense)); 17760be0d8bdSHansol Suh PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", MatMultHermitianTransposeAddColumnRange_MPIDense)); 17779566063dSJacob Faibussowitsch PetscCall(PetscObjectChangeTypeName((PetscObject)mat, MATMPIDENSE)); 17783ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1779273d9f13SBarry Smith } 1780273d9f13SBarry Smith 1781209238afSKris Buschelman /*MC 1782002d173eSKris Buschelman MATDENSE - MATDENSE = "dense" - A matrix type to be used for dense matrices. 1783209238afSKris Buschelman 178411a5261eSBarry Smith This matrix type is identical to `MATSEQDENSE` when constructed with a single process communicator, 178511a5261eSBarry Smith and `MATMPIDENSE` otherwise. 1786209238afSKris Buschelman 17872ef1f0ffSBarry Smith Options Database Key: 178811a5261eSBarry Smith . -mat_type dense - sets the matrix type to `MATDENSE` during a call to `MatSetFromOptions()` 1789209238afSKris Buschelman 1790209238afSKris Buschelman Level: beginner 1791209238afSKris Buschelman 17921cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATSEQDENSE`, `MATMPIDENSE`, `MATDENSECUDA`, `MATDENSEHIP` 17936947451fSStefano Zampini M*/ 17946947451fSStefano Zampini 17955d83a8b1SBarry Smith /*@ 1796273d9f13SBarry Smith MatMPIDenseSetPreallocation - Sets the array used to store the matrix entries 1797273d9f13SBarry Smith 17988f490ea3SStefano Zampini Collective 1799273d9f13SBarry Smith 1800273d9f13SBarry Smith Input Parameters: 1801fe59aa6dSJacob Faibussowitsch + B - the matrix 18022ef1f0ffSBarry Smith - data - optional location of matrix data. Set to `NULL` for PETSc 1803273d9f13SBarry Smith to control all matrix memory allocation. 1804273d9f13SBarry Smith 18052ef1f0ffSBarry Smith Level: intermediate 18062ef1f0ffSBarry Smith 1807273d9f13SBarry Smith Notes: 18082ef1f0ffSBarry Smith The dense format is fully compatible with standard Fortran 1809273d9f13SBarry Smith storage by columns. 1810273d9f13SBarry Smith 1811273d9f13SBarry Smith The data input variable is intended primarily for Fortran programmers 1812273d9f13SBarry Smith who wish to allocate their own matrix memory space. Most users should 18132ef1f0ffSBarry Smith set `data` to `NULL`. 1814273d9f13SBarry Smith 18151cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATMPIDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()` 1816273d9f13SBarry Smith @*/ 1817d71ae5a4SJacob Faibussowitsch PetscErrorCode MatMPIDenseSetPreallocation(Mat B, PetscScalar *data) 1818d71ae5a4SJacob Faibussowitsch { 1819273d9f13SBarry Smith PetscFunctionBegin; 1820d5ea218eSStefano Zampini PetscValidHeaderSpecific(B, MAT_CLASSID, 1); 1821cac4c232SBarry Smith PetscTryMethod(B, "MatMPIDenseSetPreallocation_C", (Mat, PetscScalar *), (B, data)); 18223ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1823273d9f13SBarry Smith } 1824273d9f13SBarry Smith 1825d3042a70SBarry Smith /*@ 182611a5261eSBarry Smith MatDensePlaceArray - Allows one to replace the array in a `MATDENSE` matrix with an 1827d3042a70SBarry Smith array provided by the user. This is useful to avoid copying an array 1828d3042a70SBarry Smith into a matrix 1829d3042a70SBarry Smith 1830d3042a70SBarry Smith Not Collective 1831d3042a70SBarry Smith 1832d3042a70SBarry Smith Input Parameters: 1833d3042a70SBarry Smith + mat - the matrix 1834d3042a70SBarry Smith - array - the array in column major order 1835d3042a70SBarry Smith 18362ef1f0ffSBarry Smith Level: developer 18372ef1f0ffSBarry Smith 183811a5261eSBarry Smith Note: 183911a5261eSBarry Smith You can return to the original array with a call to `MatDenseResetArray()`. The user is responsible for freeing this array; it will not be 1840d3042a70SBarry Smith freed when the matrix is destroyed. 1841d3042a70SBarry Smith 18421cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDenseResetArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`, 184311a5261eSBarry Smith `MatDenseReplaceArray()` 1844d3042a70SBarry Smith @*/ 1845d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDensePlaceArray(Mat mat, const PetscScalar *array) 1846d71ae5a4SJacob Faibussowitsch { 1847d3042a70SBarry Smith PetscFunctionBegin; 1848d5ea218eSStefano Zampini PetscValidHeaderSpecific(mat, MAT_CLASSID, 1); 1849cac4c232SBarry Smith PetscUseMethod(mat, "MatDensePlaceArray_C", (Mat, const PetscScalar *), (mat, array)); 18509566063dSJacob Faibussowitsch PetscCall(PetscObjectStateIncrease((PetscObject)mat)); 185147d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 1852637a0070SStefano Zampini mat->offloadmask = PETSC_OFFLOAD_CPU; 1853637a0070SStefano Zampini #endif 18543ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1855d3042a70SBarry Smith } 1856d3042a70SBarry Smith 1857d3042a70SBarry Smith /*@ 185811a5261eSBarry Smith MatDenseResetArray - Resets the matrix array to that it previously had before the call to `MatDensePlaceArray()` 1859d3042a70SBarry Smith 1860d3042a70SBarry Smith Not Collective 1861d3042a70SBarry Smith 18622fe279fdSBarry Smith Input Parameter: 1863d3042a70SBarry Smith . mat - the matrix 1864d3042a70SBarry Smith 18652ef1f0ffSBarry Smith Level: developer 18662ef1f0ffSBarry Smith 186711a5261eSBarry Smith Note: 186811a5261eSBarry Smith You can only call this after a call to `MatDensePlaceArray()` 1869d3042a70SBarry Smith 18701cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDensePlaceArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()` 1871d3042a70SBarry Smith @*/ 1872d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseResetArray(Mat mat) 1873d71ae5a4SJacob Faibussowitsch { 1874d3042a70SBarry Smith PetscFunctionBegin; 1875d5ea218eSStefano Zampini PetscValidHeaderSpecific(mat, MAT_CLASSID, 1); 1876cac4c232SBarry Smith PetscUseMethod(mat, "MatDenseResetArray_C", (Mat), (mat)); 18779566063dSJacob Faibussowitsch PetscCall(PetscObjectStateIncrease((PetscObject)mat)); 18783ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1879d3042a70SBarry Smith } 1880d3042a70SBarry Smith 1881d5ea218eSStefano Zampini /*@ 1882d5ea218eSStefano Zampini MatDenseReplaceArray - Allows one to replace the array in a dense matrix with an 1883d5ea218eSStefano Zampini array provided by the user. This is useful to avoid copying an array 1884d5ea218eSStefano Zampini into a matrix 1885d5ea218eSStefano Zampini 1886d5ea218eSStefano Zampini Not Collective 1887d5ea218eSStefano Zampini 1888d5ea218eSStefano Zampini Input Parameters: 1889d5ea218eSStefano Zampini + mat - the matrix 1890d5ea218eSStefano Zampini - array - the array in column major order 1891d5ea218eSStefano Zampini 18922ef1f0ffSBarry Smith Level: developer 18932ef1f0ffSBarry Smith 189411a5261eSBarry Smith Note: 189511a5261eSBarry Smith The memory passed in MUST be obtained with `PetscMalloc()` and CANNOT be 1896d5ea218eSStefano Zampini freed by the user. It will be freed when the matrix is destroyed. 1897d5ea218eSStefano Zampini 18981cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MatDensePlaceArray()`, `MatDenseGetArray()`, `VecReplaceArray()` 1899d5ea218eSStefano Zampini @*/ 1900d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseReplaceArray(Mat mat, const PetscScalar *array) 1901d71ae5a4SJacob Faibussowitsch { 1902d5ea218eSStefano Zampini PetscFunctionBegin; 1903d5ea218eSStefano Zampini PetscValidHeaderSpecific(mat, MAT_CLASSID, 1); 1904cac4c232SBarry Smith PetscUseMethod(mat, "MatDenseReplaceArray_C", (Mat, const PetscScalar *), (mat, array)); 19059566063dSJacob Faibussowitsch PetscCall(PetscObjectStateIncrease((PetscObject)mat)); 190647d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 1907d5ea218eSStefano Zampini mat->offloadmask = PETSC_OFFLOAD_CPU; 1908d5ea218eSStefano Zampini #endif 19093ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 1910d5ea218eSStefano Zampini } 1911d5ea218eSStefano Zampini 19125d83a8b1SBarry Smith /*@ 191311a5261eSBarry Smith MatCreateDense - Creates a matrix in `MATDENSE` format. 19148965ea79SLois Curfman McInnes 1915d083f849SBarry Smith Collective 1916db81eaa0SLois Curfman McInnes 19178965ea79SLois Curfman McInnes Input Parameters: 1918db81eaa0SLois Curfman McInnes + comm - MPI communicator 19192ef1f0ffSBarry Smith . m - number of local rows (or `PETSC_DECIDE` to have calculated if `M` is given) 19202ef1f0ffSBarry Smith . n - number of local columns (or `PETSC_DECIDE` to have calculated if `N` is given) 19212ef1f0ffSBarry Smith . M - number of global rows (or `PETSC_DECIDE` to have calculated if `m` is given) 19222ef1f0ffSBarry Smith . N - number of global columns (or `PETSC_DECIDE` to have calculated if `n` is given) 19232ef1f0ffSBarry Smith - data - optional location of matrix data. Set data to `NULL` (`PETSC_NULL_SCALAR` for Fortran users) for PETSc 1924dfc5480cSLois Curfman McInnes to control all matrix memory allocation. 19258965ea79SLois Curfman McInnes 19268965ea79SLois Curfman McInnes Output Parameter: 1927477f1c0bSLois Curfman McInnes . A - the matrix 19288965ea79SLois Curfman McInnes 19292ef1f0ffSBarry Smith Level: intermediate 19302ef1f0ffSBarry Smith 1931b259b22eSLois Curfman McInnes Notes: 19322ef1f0ffSBarry Smith The dense format is fully compatible with standard Fortran 193339ddd567SLois Curfman McInnes storage by columns. 193411a5261eSBarry Smith 193511a5261eSBarry Smith Although local portions of the matrix are stored in column-major 1936002bc6daSRichard Tran Mills order, the matrix is partitioned across MPI ranks by row. 19378965ea79SLois Curfman McInnes 193818f449edSLois Curfman McInnes The data input variable is intended primarily for Fortran programmers 193918f449edSLois Curfman McInnes who wish to allocate their own matrix memory space. Most users should 19402ef1f0ffSBarry Smith set `data` to `NULL` (`PETSC_NULL_SCALAR` for Fortran users). 194118f449edSLois Curfman McInnes 19428965ea79SLois Curfman McInnes The user MUST specify either the local or global matrix dimensions 19438965ea79SLois Curfman McInnes (possibly both). 19448965ea79SLois Curfman McInnes 19451cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()` 19468965ea79SLois Curfman McInnes @*/ 1947d71ae5a4SJacob Faibussowitsch PetscErrorCode MatCreateDense(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt M, PetscInt N, PetscScalar *data, Mat *A) 1948d71ae5a4SJacob Faibussowitsch { 19493a40ed3dSBarry Smith PetscFunctionBegin; 19509566063dSJacob Faibussowitsch PetscCall(MatCreate(comm, A)); 19519566063dSJacob Faibussowitsch PetscCall(MatSetSizes(*A, m, n, M, N)); 19526f08ad05SPierre Jolivet PetscCall(MatSetType(*A, MATDENSE)); 19539566063dSJacob Faibussowitsch PetscCall(MatSeqDenseSetPreallocation(*A, data)); 19546f08ad05SPierre Jolivet PetscCall(MatMPIDenseSetPreallocation(*A, data)); 19553ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 19568965ea79SLois Curfman McInnes } 19578965ea79SLois Curfman McInnes 1958d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDuplicate_MPIDense(Mat A, MatDuplicateOption cpvalues, Mat *newmat) 1959d71ae5a4SJacob Faibussowitsch { 19608965ea79SLois Curfman McInnes Mat mat; 19613501a2bdSLois Curfman McInnes Mat_MPIDense *a, *oldmat = (Mat_MPIDense *)A->data; 19628965ea79SLois Curfman McInnes 19633a40ed3dSBarry Smith PetscFunctionBegin; 1964f4259b30SLisandro Dalcin *newmat = NULL; 19659566063dSJacob Faibussowitsch PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat)); 19669566063dSJacob Faibussowitsch PetscCall(MatSetSizes(mat, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N)); 19679566063dSJacob Faibussowitsch PetscCall(MatSetType(mat, ((PetscObject)A)->type_name)); 1968834f8fabSBarry Smith a = (Mat_MPIDense *)mat->data; 19695aa7edbeSHong Zhang 1970d5f3da31SBarry Smith mat->factortype = A->factortype; 1971c456f294SBarry Smith mat->assembled = PETSC_TRUE; 1972273d9f13SBarry Smith mat->preallocated = PETSC_TRUE; 19738965ea79SLois Curfman McInnes 1974e0fa3b82SLois Curfman McInnes mat->insertmode = NOT_SET_VALUES; 19753782ba37SSatish Balay a->donotstash = oldmat->donotstash; 1976e04c1aa4SHong Zhang 19779566063dSJacob Faibussowitsch PetscCall(PetscLayoutReference(A->rmap, &mat->rmap)); 19789566063dSJacob Faibussowitsch PetscCall(PetscLayoutReference(A->cmap, &mat->cmap)); 19798965ea79SLois Curfman McInnes 19809566063dSJacob Faibussowitsch PetscCall(MatDuplicate(oldmat->A, cpvalues, &a->A)); 198101b82886SBarry Smith 19828965ea79SLois Curfman McInnes *newmat = mat; 19833ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 19848965ea79SLois Curfman McInnes } 19858965ea79SLois Curfman McInnes 198666976f2fSJacob Faibussowitsch static PetscErrorCode MatLoad_MPIDense(Mat newMat, PetscViewer viewer) 1987d71ae5a4SJacob Faibussowitsch { 198887d5ce66SSatish Balay PetscBool isbinary; 198987d5ce66SSatish Balay #if defined(PETSC_HAVE_HDF5) 199087d5ce66SSatish Balay PetscBool ishdf5; 199187d5ce66SSatish Balay #endif 1992eb91f321SVaclav Hapla 1993eb91f321SVaclav Hapla PetscFunctionBegin; 1994eb91f321SVaclav Hapla PetscValidHeaderSpecific(newMat, MAT_CLASSID, 1); 1995eb91f321SVaclav Hapla PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2); 1996eb91f321SVaclav Hapla /* force binary viewer to load .info file if it has not yet done so */ 19979566063dSJacob Faibussowitsch PetscCall(PetscViewerSetUp(viewer)); 19989566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary)); 199987d5ce66SSatish Balay #if defined(PETSC_HAVE_HDF5) 20009566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERHDF5, &ishdf5)); 200187d5ce66SSatish Balay #endif 2002eb91f321SVaclav Hapla if (isbinary) { 20039566063dSJacob Faibussowitsch PetscCall(MatLoad_Dense_Binary(newMat, viewer)); 2004eb91f321SVaclav Hapla #if defined(PETSC_HAVE_HDF5) 200587d5ce66SSatish Balay } else if (ishdf5) { 20069566063dSJacob Faibussowitsch PetscCall(MatLoad_Dense_HDF5(newMat, viewer)); 2007eb91f321SVaclav Hapla #endif 200898921bdaSJacob Faibussowitsch } else SETERRQ(PetscObjectComm((PetscObject)newMat), PETSC_ERR_SUP, "Viewer type %s not yet supported for reading %s matrices", ((PetscObject)viewer)->type_name, ((PetscObject)newMat)->type_name); 20093ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2010eb91f321SVaclav Hapla } 2011eb91f321SVaclav Hapla 2012d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatEqual_MPIDense(Mat A, Mat B, PetscBool *flag) 2013d71ae5a4SJacob Faibussowitsch { 20146e4ee0c6SHong Zhang Mat_MPIDense *matB = (Mat_MPIDense *)B->data, *matA = (Mat_MPIDense *)A->data; 20156e4ee0c6SHong Zhang Mat a, b; 201690ace30eSBarry Smith 20176e4ee0c6SHong Zhang PetscFunctionBegin; 20186e4ee0c6SHong Zhang a = matA->A; 20196e4ee0c6SHong Zhang b = matB->A; 20202cf15c64SPierre Jolivet PetscCall(MatEqual(a, b, flag)); 20212cf15c64SPierre Jolivet PetscCall(MPIU_Allreduce(MPI_IN_PLACE, flag, 1, MPIU_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A))); 20223ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 20236e4ee0c6SHong Zhang } 202490ace30eSBarry Smith 202566976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MatTransMatMult_MPIDense_MPIDense(void *data) 2026d71ae5a4SJacob Faibussowitsch { 20276718818eSStefano Zampini Mat_TransMatMultDense *atb = (Mat_TransMatMultDense *)data; 2028baa3c1c6SHong Zhang 2029baa3c1c6SHong Zhang PetscFunctionBegin; 20309566063dSJacob Faibussowitsch PetscCall(PetscFree2(atb->sendbuf, atb->recvcounts)); 20319566063dSJacob Faibussowitsch PetscCall(MatDestroy(&atb->atb)); 20329566063dSJacob Faibussowitsch PetscCall(PetscFree(atb)); 20333ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2034baa3c1c6SHong Zhang } 2035baa3c1c6SHong Zhang 203666976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MatMatTransMult_MPIDense_MPIDense(void *data) 2037d71ae5a4SJacob Faibussowitsch { 20386718818eSStefano Zampini Mat_MatTransMultDense *abt = (Mat_MatTransMultDense *)data; 2039cc48ffa7SToby Isaac 2040cc48ffa7SToby Isaac PetscFunctionBegin; 20419566063dSJacob Faibussowitsch PetscCall(PetscFree2(abt->buf[0], abt->buf[1])); 20429566063dSJacob Faibussowitsch PetscCall(PetscFree2(abt->recvcounts, abt->recvdispls)); 20439566063dSJacob Faibussowitsch PetscCall(PetscFree(abt)); 20443ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2045cc48ffa7SToby Isaac } 2046cc48ffa7SToby Isaac 2047d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C) 2048d71ae5a4SJacob Faibussowitsch { 2049baa3c1c6SHong Zhang Mat_MPIDense *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data; 20506718818eSStefano Zampini Mat_TransMatMultDense *atb; 2051cb20be35SHong Zhang MPI_Comm comm; 20526718818eSStefano Zampini PetscMPIInt size, *recvcounts; 20536718818eSStefano Zampini PetscScalar *carray, *sendbuf; 2054637a0070SStefano Zampini const PetscScalar *atbarray; 2055a1f9a5e6SStefano Zampini PetscInt i, cN = C->cmap->N, proc, k, j, lda; 2056e68c0b26SHong Zhang const PetscInt *ranges; 2057cb20be35SHong Zhang 2058cb20be35SHong Zhang PetscFunctionBegin; 20596718818eSStefano Zampini MatCheckProduct(C, 3); 20605f80ce2aSJacob Faibussowitsch PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty"); 20616718818eSStefano Zampini atb = (Mat_TransMatMultDense *)C->product->data; 20626718818eSStefano Zampini recvcounts = atb->recvcounts; 20636718818eSStefano Zampini sendbuf = atb->sendbuf; 20646718818eSStefano Zampini 20659566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A, &comm)); 20669566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm, &size)); 2067e68c0b26SHong Zhang 2068c5ef1628SHong Zhang /* compute atbarray = aseq^T * bseq */ 2069fb842aefSJose E. Roman PetscCall(MatTransposeMatMult(a->A, b->A, atb->atb ? MAT_REUSE_MATRIX : MAT_INITIAL_MATRIX, PETSC_DETERMINE, &atb->atb)); 2070cb20be35SHong Zhang 20719566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRanges(C, &ranges)); 2072cb20be35SHong Zhang 2073660d5466SHong Zhang /* arrange atbarray into sendbuf */ 20749566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(atb->atb, &atbarray)); 20759566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(atb->atb, &lda)); 2076637a0070SStefano Zampini for (proc = 0, k = 0; proc < size; proc++) { 2077baa3c1c6SHong Zhang for (j = 0; j < cN; j++) { 2078a1f9a5e6SStefano Zampini for (i = ranges[proc]; i < ranges[proc + 1]; i++) sendbuf[k++] = atbarray[i + j * lda]; 2079cb20be35SHong Zhang } 2080cb20be35SHong Zhang } 20819566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(atb->atb, &atbarray)); 2082637a0070SStefano Zampini 2083c5ef1628SHong Zhang /* sum all atbarray to local values of C */ 20849566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayWrite(c->A, &carray)); 20859566063dSJacob Faibussowitsch PetscCallMPI(MPI_Reduce_scatter(sendbuf, carray, recvcounts, MPIU_SCALAR, MPIU_SUM, comm)); 20869566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayWrite(c->A, &carray)); 208767af85e8SPierre Jolivet PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE)); 20889566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY)); 20899566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY)); 20903ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2091cb20be35SHong Zhang } 2092cb20be35SHong Zhang 2093d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C) 2094d71ae5a4SJacob Faibussowitsch { 2095cb20be35SHong Zhang MPI_Comm comm; 2096baa3c1c6SHong Zhang PetscMPIInt size; 2097660d5466SHong Zhang PetscInt cm = A->cmap->n, cM, cN = B->cmap->N; 2098baa3c1c6SHong Zhang Mat_TransMatMultDense *atb; 209947d993e7Ssuyashtn PetscBool cisdense = PETSC_FALSE; 21000acaeb71SStefano Zampini PetscInt i; 21010acaeb71SStefano Zampini const PetscInt *ranges; 2102cb20be35SHong Zhang 2103cb20be35SHong Zhang PetscFunctionBegin; 21048f83a4d9SJacob Faibussowitsch MatCheckProduct(C, 4); 21055f80ce2aSJacob Faibussowitsch PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty"); 21069566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A, &comm)); 2107cb20be35SHong Zhang if (A->rmap->rstart != B->rmap->rstart || A->rmap->rend != B->rmap->rend) { 210898921bdaSJacob Faibussowitsch SETERRQ(comm, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != B (%" PetscInt_FMT ",%" PetscInt_FMT ")", A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend); 2109cb20be35SHong Zhang } 2110cb20be35SHong Zhang 21114222ddf1SHong Zhang /* create matrix product C */ 21129566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C, cm, B->cmap->n, A->cmap->N, B->cmap->N)); 211347d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) 21149566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSECUDA, "")); 211547d993e7Ssuyashtn #endif 211647d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP) 211747d993e7Ssuyashtn PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSEHIP, "")); 211847d993e7Ssuyashtn #endif 211948a46eb9SPierre Jolivet if (!cisdense) PetscCall(MatSetType(C, ((PetscObject)A)->type_name)); 21209566063dSJacob Faibussowitsch PetscCall(MatSetUp(C)); 2121baa3c1c6SHong Zhang 21224222ddf1SHong Zhang /* create data structure for reuse C */ 21239566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm, &size)); 21249566063dSJacob Faibussowitsch PetscCall(PetscNew(&atb)); 21254222ddf1SHong Zhang cM = C->rmap->N; 21269566063dSJacob Faibussowitsch PetscCall(PetscMalloc2(cM * cN, &atb->sendbuf, size, &atb->recvcounts)); 21279566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRanges(C, &ranges)); 21280acaeb71SStefano Zampini for (i = 0; i < size; i++) atb->recvcounts[i] = (ranges[i + 1] - ranges[i]) * cN; 2129baa3c1c6SHong Zhang 21306718818eSStefano Zampini C->product->data = atb; 21316718818eSStefano Zampini C->product->destroy = MatDestroy_MatTransMatMult_MPIDense_MPIDense; 21323ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2133cb20be35SHong Zhang } 2134cb20be35SHong Zhang 2135d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C) 2136d71ae5a4SJacob Faibussowitsch { 2137cc48ffa7SToby Isaac MPI_Comm comm; 2138cc48ffa7SToby Isaac PetscMPIInt i, size; 2139cc48ffa7SToby Isaac PetscInt maxRows, bufsiz; 2140cc48ffa7SToby Isaac PetscMPIInt tag; 21414222ddf1SHong Zhang PetscInt alg; 2142cc48ffa7SToby Isaac Mat_MatTransMultDense *abt; 21434222ddf1SHong Zhang Mat_Product *product = C->product; 21444222ddf1SHong Zhang PetscBool flg; 2145cc48ffa7SToby Isaac 2146cc48ffa7SToby Isaac PetscFunctionBegin; 21476718818eSStefano Zampini MatCheckProduct(C, 4); 21485f80ce2aSJacob Faibussowitsch PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty"); 21494222ddf1SHong Zhang /* check local size of A and B */ 21505f80ce2aSJacob Faibussowitsch PetscCheck(A->cmap->n == B->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Matrix local column dimensions are incompatible, A (%" PetscInt_FMT ") != B (%" PetscInt_FMT ")", A->cmap->n, B->cmap->n); 2151cc48ffa7SToby Isaac 21529566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg, "allgatherv", &flg)); 2153637a0070SStefano Zampini alg = flg ? 0 : 1; 2154cc48ffa7SToby Isaac 21554222ddf1SHong Zhang /* setup matrix product C */ 21569566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C, A->rmap->n, B->rmap->n, A->rmap->N, B->rmap->N)); 21579566063dSJacob Faibussowitsch PetscCall(MatSetType(C, MATMPIDENSE)); 21589566063dSJacob Faibussowitsch PetscCall(MatSetUp(C)); 21599566063dSJacob Faibussowitsch PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag)); 21604222ddf1SHong Zhang 21614222ddf1SHong Zhang /* create data structure for reuse C */ 21629566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C, &comm)); 21639566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm, &size)); 21649566063dSJacob Faibussowitsch PetscCall(PetscNew(&abt)); 2165cc48ffa7SToby Isaac abt->tag = tag; 2166faa55883SToby Isaac abt->alg = alg; 2167faa55883SToby Isaac switch (alg) { 21684222ddf1SHong Zhang case 1: /* alg: "cyclic" */ 2169cc48ffa7SToby Isaac for (maxRows = 0, i = 0; i < size; i++) maxRows = PetscMax(maxRows, (B->rmap->range[i + 1] - B->rmap->range[i])); 2170cc48ffa7SToby Isaac bufsiz = A->cmap->N * maxRows; 2171f4f49eeaSPierre Jolivet PetscCall(PetscMalloc2(bufsiz, &abt->buf[0], bufsiz, &abt->buf[1])); 2172faa55883SToby Isaac break; 21734222ddf1SHong Zhang default: /* alg: "allgatherv" */ 2174f4f49eeaSPierre Jolivet PetscCall(PetscMalloc2(B->rmap->n * B->cmap->N, &abt->buf[0], B->rmap->N * B->cmap->N, &abt->buf[1])); 2175f4f49eeaSPierre Jolivet PetscCall(PetscMalloc2(size, &abt->recvcounts, size + 1, &abt->recvdispls)); 2176faa55883SToby Isaac for (i = 0; i <= size; i++) abt->recvdispls[i] = B->rmap->range[i] * A->cmap->N; 2177faa55883SToby Isaac for (i = 0; i < size; i++) abt->recvcounts[i] = abt->recvdispls[i + 1] - abt->recvdispls[i]; 2178faa55883SToby Isaac break; 2179faa55883SToby Isaac } 2180cc48ffa7SToby Isaac 21816718818eSStefano Zampini C->product->data = abt; 21826718818eSStefano Zampini C->product->destroy = MatDestroy_MatMatTransMult_MPIDense_MPIDense; 21834222ddf1SHong Zhang C->ops->mattransposemultnumeric = MatMatTransposeMultNumeric_MPIDense_MPIDense; 21843ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2185cc48ffa7SToby Isaac } 2186cc48ffa7SToby Isaac 2187d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(Mat A, Mat B, Mat C) 2188d71ae5a4SJacob Faibussowitsch { 2189cc48ffa7SToby Isaac Mat_MPIDense *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data; 21906718818eSStefano Zampini Mat_MatTransMultDense *abt; 2191cc48ffa7SToby Isaac MPI_Comm comm; 2192cc48ffa7SToby Isaac PetscMPIInt rank, size, sendsiz, recvsiz, sendto, recvfrom, recvisfrom; 2193f4259b30SLisandro Dalcin PetscScalar *sendbuf, *recvbuf = NULL, *cv; 2194cc48ffa7SToby Isaac PetscInt i, cK = A->cmap->N, k, j, bn; 2195cc48ffa7SToby Isaac PetscScalar _DOne = 1.0, _DZero = 0.0; 2196637a0070SStefano Zampini const PetscScalar *av, *bv; 219723fff9afSBarry Smith PetscBLASInt cm, cn, ck, alda, blda = 0, clda; 2198cc48ffa7SToby Isaac MPI_Request reqs[2]; 2199cc48ffa7SToby Isaac const PetscInt *ranges; 2200cc48ffa7SToby Isaac 2201cc48ffa7SToby Isaac PetscFunctionBegin; 22026718818eSStefano Zampini MatCheckProduct(C, 3); 22035f80ce2aSJacob Faibussowitsch PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty"); 22046718818eSStefano Zampini abt = (Mat_MatTransMultDense *)C->product->data; 22059566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C, &comm)); 22069566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm, &rank)); 22079566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm, &size)); 22089566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(a->A, &av)); 22099566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(b->A, &bv)); 22109566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayWrite(c->A, &cv)); 22119566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &i)); 22129566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(i, &alda)); 22139566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(b->A, &i)); 22149566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(i, &blda)); 22159566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(c->A, &i)); 22169566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(i, &clda)); 22179566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRanges(B, &ranges)); 2218cc48ffa7SToby Isaac bn = B->rmap->n; 2219637a0070SStefano Zampini if (blda == bn) { 2220637a0070SStefano Zampini sendbuf = (PetscScalar *)bv; 2221cc48ffa7SToby Isaac } else { 2222cc48ffa7SToby Isaac sendbuf = abt->buf[0]; 2223cc48ffa7SToby Isaac for (k = 0, i = 0; i < cK; i++) { 2224ad540459SPierre Jolivet for (j = 0; j < bn; j++, k++) sendbuf[k] = bv[i * blda + j]; 2225cc48ffa7SToby Isaac } 2226cc48ffa7SToby Isaac } 2227cc48ffa7SToby Isaac if (size > 1) { 2228cc48ffa7SToby Isaac sendto = (rank + size - 1) % size; 2229cc48ffa7SToby Isaac recvfrom = (rank + size + 1) % size; 2230cc48ffa7SToby Isaac } else { 2231cc48ffa7SToby Isaac sendto = recvfrom = 0; 2232cc48ffa7SToby Isaac } 22339566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(cK, &ck)); 22349566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm)); 2235cc48ffa7SToby Isaac recvisfrom = rank; 2236cc48ffa7SToby Isaac for (i = 0; i < size; i++) { 2237cc48ffa7SToby Isaac /* we have finished receiving in sending, bufs can be read/modified */ 2238cc48ffa7SToby Isaac PetscInt nextrecvisfrom = (recvisfrom + 1) % size; /* which process the next recvbuf will originate on */ 2239cc48ffa7SToby Isaac PetscInt nextbn = ranges[nextrecvisfrom + 1] - ranges[nextrecvisfrom]; 2240cc48ffa7SToby Isaac 2241cc48ffa7SToby Isaac if (nextrecvisfrom != rank) { 2242cc48ffa7SToby Isaac /* start the cyclic sends from sendbuf, to recvbuf (which will switch to sendbuf) */ 2243cc48ffa7SToby Isaac sendsiz = cK * bn; 2244cc48ffa7SToby Isaac recvsiz = cK * nextbn; 2245cc48ffa7SToby Isaac recvbuf = (i & 1) ? abt->buf[0] : abt->buf[1]; 22469566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(sendbuf, sendsiz, MPIU_SCALAR, sendto, abt->tag, comm, &reqs[0])); 22479566063dSJacob Faibussowitsch PetscCallMPI(MPI_Irecv(recvbuf, recvsiz, MPIU_SCALAR, recvfrom, abt->tag, comm, &reqs[1])); 2248cc48ffa7SToby Isaac } 2249cc48ffa7SToby Isaac 2250cc48ffa7SToby Isaac /* local aseq * sendbuf^T */ 22519566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(ranges[recvisfrom + 1] - ranges[recvisfrom], &cn)); 2252792fecdfSBarry Smith if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "T", &cm, &cn, &ck, &_DOne, av, &alda, sendbuf, &cn, &_DZero, cv + clda * ranges[recvisfrom], &clda)); 2253cc48ffa7SToby Isaac 2254cc48ffa7SToby Isaac if (nextrecvisfrom != rank) { 2255cc48ffa7SToby Isaac /* wait for the sends and receives to complete, swap sendbuf and recvbuf */ 22569566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitall(2, reqs, MPI_STATUSES_IGNORE)); 2257cc48ffa7SToby Isaac } 2258cc48ffa7SToby Isaac bn = nextbn; 2259cc48ffa7SToby Isaac recvisfrom = nextrecvisfrom; 2260cc48ffa7SToby Isaac sendbuf = recvbuf; 2261cc48ffa7SToby Isaac } 22629566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(a->A, &av)); 22639566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(b->A, &bv)); 22649566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayWrite(c->A, &cv)); 226567af85e8SPierre Jolivet PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE)); 22669566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY)); 22679566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY)); 22683ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2269cc48ffa7SToby Isaac } 2270cc48ffa7SToby Isaac 2271d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(Mat A, Mat B, Mat C) 2272d71ae5a4SJacob Faibussowitsch { 2273faa55883SToby Isaac Mat_MPIDense *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data; 22746718818eSStefano Zampini Mat_MatTransMultDense *abt; 2275faa55883SToby Isaac MPI_Comm comm; 2276637a0070SStefano Zampini PetscMPIInt size; 2277637a0070SStefano Zampini PetscScalar *cv, *sendbuf, *recvbuf; 2278637a0070SStefano Zampini const PetscScalar *av, *bv; 2279637a0070SStefano Zampini PetscInt blda, i, cK = A->cmap->N, k, j, bn; 2280faa55883SToby Isaac PetscScalar _DOne = 1.0, _DZero = 0.0; 2281637a0070SStefano Zampini PetscBLASInt cm, cn, ck, alda, clda; 2282faa55883SToby Isaac 2283faa55883SToby Isaac PetscFunctionBegin; 22846718818eSStefano Zampini MatCheckProduct(C, 3); 22855f80ce2aSJacob Faibussowitsch PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty"); 22866718818eSStefano Zampini abt = (Mat_MatTransMultDense *)C->product->data; 22879566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A, &comm)); 22889566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm, &size)); 22899566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(a->A, &av)); 22909566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(b->A, &bv)); 22919566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayWrite(c->A, &cv)); 22929566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &i)); 22939566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(i, &alda)); 22949566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(b->A, &blda)); 22959566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(c->A, &i)); 22969566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(i, &clda)); 2297faa55883SToby Isaac /* copy transpose of B into buf[0] */ 2298faa55883SToby Isaac bn = B->rmap->n; 2299faa55883SToby Isaac sendbuf = abt->buf[0]; 2300faa55883SToby Isaac recvbuf = abt->buf[1]; 2301faa55883SToby Isaac for (k = 0, j = 0; j < bn; j++) { 2302ad540459SPierre Jolivet for (i = 0; i < cK; i++, k++) sendbuf[k] = bv[i * blda + j]; 2303faa55883SToby Isaac } 23049566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(b->A, &bv)); 23059566063dSJacob Faibussowitsch PetscCallMPI(MPI_Allgatherv(sendbuf, bn * cK, MPIU_SCALAR, recvbuf, abt->recvcounts, abt->recvdispls, MPIU_SCALAR, comm)); 23069566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(cK, &ck)); 23079566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm)); 23089566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(c->A->cmap->n, &cn)); 2309792fecdfSBarry Smith if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "N", &cm, &cn, &ck, &_DOne, av, &alda, recvbuf, &ck, &_DZero, cv, &clda)); 23109566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(a->A, &av)); 23119566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(b->A, &bv)); 23129566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayWrite(c->A, &cv)); 231367af85e8SPierre Jolivet PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE)); 23149566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY)); 23159566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY)); 23163ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2317faa55883SToby Isaac } 2318faa55883SToby Isaac 2319d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C) 2320d71ae5a4SJacob Faibussowitsch { 23216718818eSStefano Zampini Mat_MatTransMultDense *abt; 2322faa55883SToby Isaac 2323faa55883SToby Isaac PetscFunctionBegin; 23246718818eSStefano Zampini MatCheckProduct(C, 3); 23255f80ce2aSJacob Faibussowitsch PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty"); 23266718818eSStefano Zampini abt = (Mat_MatTransMultDense *)C->product->data; 2327faa55883SToby Isaac switch (abt->alg) { 2328d71ae5a4SJacob Faibussowitsch case 1: 2329d71ae5a4SJacob Faibussowitsch PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(A, B, C)); 2330d71ae5a4SJacob Faibussowitsch break; 2331d71ae5a4SJacob Faibussowitsch default: 2332d71ae5a4SJacob Faibussowitsch PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(A, B, C)); 2333d71ae5a4SJacob Faibussowitsch break; 2334faa55883SToby Isaac } 23353ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2336faa55883SToby Isaac } 2337faa55883SToby Isaac 233866976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MatMatMult_MPIDense_MPIDense(void *data) 2339d71ae5a4SJacob Faibussowitsch { 23406718818eSStefano Zampini Mat_MatMultDense *ab = (Mat_MatMultDense *)data; 2341320f2790SHong Zhang 2342320f2790SHong Zhang PetscFunctionBegin; 23439566063dSJacob Faibussowitsch PetscCall(MatDestroy(&ab->Ce)); 23449566063dSJacob Faibussowitsch PetscCall(MatDestroy(&ab->Ae)); 23459566063dSJacob Faibussowitsch PetscCall(MatDestroy(&ab->Be)); 23469566063dSJacob Faibussowitsch PetscCall(PetscFree(ab)); 23473ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2348320f2790SHong Zhang } 2349320f2790SHong Zhang 23505d8c7819SPierre Jolivet static PetscErrorCode MatMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C) 2351d71ae5a4SJacob Faibussowitsch { 23526718818eSStefano Zampini Mat_MatMultDense *ab; 23535d8c7819SPierre Jolivet Mat_MPIDense *mdn = (Mat_MPIDense *)A->data; 2354320f2790SHong Zhang 2355320f2790SHong Zhang PetscFunctionBegin; 23566718818eSStefano Zampini MatCheckProduct(C, 3); 23575f80ce2aSJacob Faibussowitsch PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Missing product data"); 23586718818eSStefano Zampini ab = (Mat_MatMultDense *)C->product->data; 23595d8c7819SPierre Jolivet if (ab->Ae && ab->Ce) { 23605d8c7819SPierre Jolivet #if PetscDefined(HAVE_ELEMENTAL) 23619566063dSJacob Faibussowitsch PetscCall(MatConvert_MPIDense_Elemental(A, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Ae)); 23629566063dSJacob Faibussowitsch PetscCall(MatConvert_MPIDense_Elemental(B, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Be)); 23639566063dSJacob Faibussowitsch PetscCall(MatMatMultNumeric_Elemental(ab->Ae, ab->Be, ab->Ce)); 23649566063dSJacob Faibussowitsch PetscCall(MatConvert(ab->Ce, MATMPIDENSE, MAT_REUSE_MATRIX, &C)); 23655d8c7819SPierre Jolivet #else 23665d8c7819SPierre Jolivet SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined"); 23675d8c7819SPierre Jolivet #endif 23685d8c7819SPierre Jolivet } else { 23695d8c7819SPierre Jolivet const PetscScalar *read; 23705d8c7819SPierre Jolivet PetscScalar *write; 23715d8c7819SPierre Jolivet PetscInt lda; 23725d8c7819SPierre Jolivet 23735d8c7819SPierre Jolivet PetscCall(MatDenseGetLDA(B, &lda)); 23745d8c7819SPierre Jolivet PetscCall(MatDenseGetArrayRead(B, &read)); 23755d8c7819SPierre Jolivet PetscCall(MatDenseGetArrayWrite(ab->Be, &write)); 23765d8c7819SPierre Jolivet if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A)); /* cannot be done during the symbolic phase because of possible calls to MatProductReplaceMats() */ 23775d8c7819SPierre Jolivet for (PetscInt i = 0; i < C->cmap->N; ++i) { 23785d8c7819SPierre Jolivet PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE)); 23795d8c7819SPierre Jolivet PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE)); 23805d8c7819SPierre Jolivet } 23815d8c7819SPierre Jolivet PetscCall(MatDenseRestoreArrayWrite(ab->Be, &write)); 23825d8c7819SPierre Jolivet PetscCall(MatDenseRestoreArrayRead(B, &read)); 2383f4f49eeaSPierre Jolivet PetscCall(MatMatMultNumeric_SeqDense_SeqDense(((Mat_MPIDense *)A->data)->A, ab->Be, ((Mat_MPIDense *)C->data)->A)); 23845d8c7819SPierre Jolivet } 23853ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2386320f2790SHong Zhang } 2387320f2790SHong Zhang 2388d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C) 2389d71ae5a4SJacob Faibussowitsch { 23905d8c7819SPierre Jolivet Mat_Product *product = C->product; 23915d8c7819SPierre Jolivet PetscInt alg; 2392320f2790SHong Zhang Mat_MatMultDense *ab; 23935d8c7819SPierre Jolivet PetscBool flg; 2394320f2790SHong Zhang 2395320f2790SHong Zhang PetscFunctionBegin; 23966718818eSStefano Zampini MatCheckProduct(C, 4); 23975f80ce2aSJacob Faibussowitsch PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty"); 23984222ddf1SHong Zhang /* check local size of A and B */ 23995d8c7819SPierre Jolivet PetscCheck(A->cmap->rstart == B->rmap->rstart && A->cmap->rend == B->rmap->rend, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != B (%" PetscInt_FMT ", %" PetscInt_FMT ")", 24005d8c7819SPierre Jolivet A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend); 2401320f2790SHong Zhang 24025d8c7819SPierre Jolivet PetscCall(PetscStrcmp(product->alg, "petsc", &flg)); 24035d8c7819SPierre Jolivet alg = flg ? 0 : 1; 24044222ddf1SHong Zhang 24054222ddf1SHong Zhang /* setup C */ 24065d8c7819SPierre Jolivet PetscCall(MatSetSizes(C, A->rmap->n, B->cmap->n, A->rmap->N, B->cmap->N)); 24075d8c7819SPierre Jolivet PetscCall(MatSetType(C, MATMPIDENSE)); 24089566063dSJacob Faibussowitsch PetscCall(MatSetUp(C)); 2409320f2790SHong Zhang 2410320f2790SHong Zhang /* create data structure for reuse Cdense */ 24119566063dSJacob Faibussowitsch PetscCall(PetscNew(&ab)); 24125d8c7819SPierre Jolivet 24135d8c7819SPierre Jolivet switch (alg) { 24145d8c7819SPierre Jolivet case 1: /* alg: "elemental" */ 24155d8c7819SPierre Jolivet #if PetscDefined(HAVE_ELEMENTAL) 24165d8c7819SPierre Jolivet /* create elemental matrices Ae and Be */ 24175d8c7819SPierre Jolivet PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &ab->Ae)); 24185d8c7819SPierre Jolivet PetscCall(MatSetSizes(ab->Ae, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N)); 24195d8c7819SPierre Jolivet PetscCall(MatSetType(ab->Ae, MATELEMENTAL)); 24205d8c7819SPierre Jolivet PetscCall(MatSetUp(ab->Ae)); 24215d8c7819SPierre Jolivet PetscCall(MatSetOption(ab->Ae, MAT_ROW_ORIENTED, PETSC_FALSE)); 24225d8c7819SPierre Jolivet 24235d8c7819SPierre Jolivet PetscCall(MatCreate(PetscObjectComm((PetscObject)B), &ab->Be)); 24245d8c7819SPierre Jolivet PetscCall(MatSetSizes(ab->Be, PETSC_DECIDE, PETSC_DECIDE, B->rmap->N, B->cmap->N)); 24255d8c7819SPierre Jolivet PetscCall(MatSetType(ab->Be, MATELEMENTAL)); 24265d8c7819SPierre Jolivet PetscCall(MatSetUp(ab->Be)); 24275d8c7819SPierre Jolivet PetscCall(MatSetOption(ab->Be, MAT_ROW_ORIENTED, PETSC_FALSE)); 24285d8c7819SPierre Jolivet 24295d8c7819SPierre Jolivet /* compute symbolic Ce = Ae*Be */ 24305d8c7819SPierre Jolivet PetscCall(MatCreate(PetscObjectComm((PetscObject)C), &ab->Ce)); 24315d8c7819SPierre Jolivet PetscCall(MatMatMultSymbolic_Elemental(ab->Ae, ab->Be, fill, ab->Ce)); 24325d8c7819SPierre Jolivet #else 24335d8c7819SPierre Jolivet SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined"); 24345d8c7819SPierre Jolivet #endif 24355d8c7819SPierre Jolivet break; 24365d8c7819SPierre Jolivet default: /* alg: "petsc" */ 24375d8c7819SPierre Jolivet ab->Ae = NULL; 24385d8c7819SPierre Jolivet PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, A->cmap->N, B->cmap->N, NULL, &ab->Be)); 24395d8c7819SPierre Jolivet ab->Ce = NULL; 24405d8c7819SPierre Jolivet break; 24415d8c7819SPierre Jolivet } 24426718818eSStefano Zampini 24436718818eSStefano Zampini C->product->data = ab; 24446718818eSStefano Zampini C->product->destroy = MatDestroy_MatMatMult_MPIDense_MPIDense; 24454222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIDense_MPIDense; 24463ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2447320f2790SHong Zhang } 24482ef1f0ffSBarry Smith 2449d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_AB(Mat C) 2450d71ae5a4SJacob Faibussowitsch { 24515d8c7819SPierre Jolivet Mat_Product *product = C->product; 24525d8c7819SPierre Jolivet const char *algTypes[2] = {"petsc", "elemental"}; 24535d8c7819SPierre Jolivet PetscInt alg, nalg = PetscDefined(HAVE_ELEMENTAL) ? 2 : 1; 24545d8c7819SPierre Jolivet PetscBool flg = PETSC_FALSE; 24555d8c7819SPierre Jolivet 24564d86920dSPierre Jolivet PetscFunctionBegin; 24575d8c7819SPierre Jolivet /* Set default algorithm */ 24585d8c7819SPierre Jolivet alg = 0; /* default is petsc */ 24595d8c7819SPierre Jolivet PetscCall(PetscStrcmp(product->alg, "default", &flg)); 24605d8c7819SPierre Jolivet if (flg) PetscCall(MatProductSetAlgorithm(C, (MatProductAlgorithm)algTypes[alg])); 24615d8c7819SPierre Jolivet 24625d8c7819SPierre Jolivet /* Get runtime option */ 24635d8c7819SPierre Jolivet PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_AB", "Mat"); 24645d8c7819SPierre Jolivet PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_AB", algTypes, nalg, algTypes[alg], &alg, &flg)); 24655d8c7819SPierre Jolivet PetscOptionsEnd(); 24665d8c7819SPierre Jolivet if (flg) PetscCall(MatProductSetAlgorithm(C, (MatProductAlgorithm)algTypes[alg])); 24675d8c7819SPierre Jolivet 24684222ddf1SHong Zhang C->ops->matmultsymbolic = MatMatMultSymbolic_MPIDense_MPIDense; 24694222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AB; 24703ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 2471320f2790SHong Zhang } 247286aefd0dSHong Zhang 2473d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_AtB(Mat C) 2474d71ae5a4SJacob Faibussowitsch { 24754222ddf1SHong Zhang Mat_Product *product = C->product; 24764222ddf1SHong Zhang Mat A = product->A, B = product->B; 24774222ddf1SHong Zhang 24784222ddf1SHong Zhang PetscFunctionBegin; 247907e90273SPierre Jolivet PetscCheck(A->rmap->rstart == B->rmap->rstart && A->rmap->rend == B->rmap->rend, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, (%" PetscInt_FMT ", %" PetscInt_FMT ") != (%" PetscInt_FMT ",%" PetscInt_FMT ")", 248007e90273SPierre Jolivet A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend); 24816718818eSStefano Zampini C->ops->transposematmultsymbolic = MatTransposeMatMultSymbolic_MPIDense_MPIDense; 24826718818eSStefano Zampini C->ops->productsymbolic = MatProductSymbolic_AtB; 24833ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 24844222ddf1SHong Zhang } 24854222ddf1SHong Zhang 2486d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_ABt(Mat C) 2487d71ae5a4SJacob Faibussowitsch { 24884222ddf1SHong Zhang Mat_Product *product = C->product; 24894222ddf1SHong Zhang const char *algTypes[2] = {"allgatherv", "cyclic"}; 24904222ddf1SHong Zhang PetscInt alg, nalg = 2; 24914222ddf1SHong Zhang PetscBool flg = PETSC_FALSE; 24924222ddf1SHong Zhang 24934222ddf1SHong Zhang PetscFunctionBegin; 24944222ddf1SHong Zhang /* Set default algorithm */ 24954222ddf1SHong Zhang alg = 0; /* default is allgatherv */ 24969566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg, "default", &flg)); 249748a46eb9SPierre Jolivet if (flg) PetscCall(MatProductSetAlgorithm(C, (MatProductAlgorithm)algTypes[alg])); 24984222ddf1SHong Zhang 24994222ddf1SHong Zhang /* Get runtime option */ 25004222ddf1SHong Zhang if (product->api_user) { 2501d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatMatTransposeMult", "Mat"); 25029566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matmattransmult_mpidense_mpidense_via", "Algorithmic approach", "MatMatTransposeMult", algTypes, nalg, algTypes[alg], &alg, &flg)); 2503d0609cedSBarry Smith PetscOptionsEnd(); 25044222ddf1SHong Zhang } else { 2505d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_ABt", "Mat"); 25069566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_ABt", algTypes, nalg, algTypes[alg], &alg, &flg)); 2507d0609cedSBarry Smith PetscOptionsEnd(); 25084222ddf1SHong Zhang } 250948a46eb9SPierre Jolivet if (flg) PetscCall(MatProductSetAlgorithm(C, (MatProductAlgorithm)algTypes[alg])); 25104222ddf1SHong Zhang 25114222ddf1SHong Zhang C->ops->mattransposemultsymbolic = MatMatTransposeMultSymbolic_MPIDense_MPIDense; 25124222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_ABt; 25133ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 25144222ddf1SHong Zhang } 25154222ddf1SHong Zhang 25165d8c7819SPierre Jolivet static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat C) 2517d71ae5a4SJacob Faibussowitsch { 25184222ddf1SHong Zhang Mat_Product *product = C->product; 25194222ddf1SHong Zhang 25204222ddf1SHong Zhang PetscFunctionBegin; 25214222ddf1SHong Zhang switch (product->type) { 2522d71ae5a4SJacob Faibussowitsch case MATPRODUCT_AB: 2523d71ae5a4SJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIDense_AB(C)); 2524d71ae5a4SJacob Faibussowitsch break; 2525d71ae5a4SJacob Faibussowitsch case MATPRODUCT_AtB: 2526d71ae5a4SJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIDense_AtB(C)); 2527d71ae5a4SJacob Faibussowitsch break; 2528d71ae5a4SJacob Faibussowitsch case MATPRODUCT_ABt: 2529d71ae5a4SJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIDense_ABt(C)); 2530d71ae5a4SJacob Faibussowitsch break; 2531d71ae5a4SJacob Faibussowitsch default: 2532d71ae5a4SJacob Faibussowitsch break; 25334222ddf1SHong Zhang } 25343ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 25354222ddf1SHong Zhang } 2536