xref: /petsc/src/mat/impls/dense/mpi/mpidense.c (revision 4d5500e8ab0f8e7b1fe2746469a330fafa8aaa90)
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) {
883462c564dSBarry Smith     PetscCallMPI(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) {
891462c564dSBarry Smith     PetscCallMPI(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:
9269725ddffSPierre Jolivet   case MAT_IGNORE_LOWER_TRIANGULAR:
9279725ddffSPierre Jolivet   case MAT_IGNORE_ZERO_ENTRIES:
9289725ddffSPierre Jolivet   case MAT_SUBMAT_SINGLEIS:
929d71ae5a4SJacob Faibussowitsch     PetscCall(PetscInfo(A, "Option %s ignored\n", MatOptions[op]));
930d71ae5a4SJacob Faibussowitsch     break;
931d71ae5a4SJacob Faibussowitsch   case MAT_IGNORE_OFF_PROC_ENTRIES:
932d71ae5a4SJacob Faibussowitsch     a->donotstash = flg;
933d71ae5a4SJacob Faibussowitsch     break;
93477e54ba9SKris Buschelman   case MAT_SYMMETRIC:
93577e54ba9SKris Buschelman   case MAT_STRUCTURALLY_SYMMETRIC:
9369a4540c5SBarry Smith   case MAT_HERMITIAN:
9379a4540c5SBarry Smith   case MAT_SYMMETRY_ETERNAL:
938b94d7dedSBarry Smith   case MAT_STRUCTURAL_SYMMETRY_ETERNAL:
939b94d7dedSBarry Smith   case MAT_SPD:
940b94d7dedSBarry Smith   case MAT_SPD_ETERNAL:
941b94d7dedSBarry Smith     /* if the diagonal matrix is square it inherits some of the properties above */
9429725ddffSPierre Jolivet     if (a->A && A->rmap->n == A->cmap->n) PetscCall(MatSetOption(a->A, op, flg));
94377e54ba9SKris Buschelman     break;
944d71ae5a4SJacob Faibussowitsch   default:
945d71ae5a4SJacob Faibussowitsch     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "unknown option %s", MatOptions[op]);
9463a40ed3dSBarry Smith   }
9473ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9488965ea79SLois Curfman McInnes }
9498965ea79SLois Curfman McInnes 
95066976f2fSJacob Faibussowitsch static PetscErrorCode MatDiagonalScale_MPIDense(Mat A, Vec ll, Vec rr)
951d71ae5a4SJacob Faibussowitsch {
9525b2fa520SLois Curfman McInnes   Mat_MPIDense      *mdn = (Mat_MPIDense *)A->data;
953637a0070SStefano Zampini   const PetscScalar *l;
954637a0070SStefano Zampini   PetscScalar        x, *v, *vv, *r;
955637a0070SStefano Zampini   PetscInt           i, j, s2a, s3a, s2, s3, m = mdn->A->rmap->n, n = mdn->A->cmap->n, lda;
9565b2fa520SLois Curfman McInnes 
9575b2fa520SLois Curfman McInnes   PetscFunctionBegin;
9589566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(mdn->A, &vv));
9599566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(mdn->A, &lda));
9609566063dSJacob Faibussowitsch   PetscCall(MatGetLocalSize(A, &s2, &s3));
9615b2fa520SLois Curfman McInnes   if (ll) {
9629566063dSJacob Faibussowitsch     PetscCall(VecGetLocalSize(ll, &s2a));
96308401ef6SPierre Jolivet     PetscCheck(s2a == s2, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Left scaling vector non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT, s2a, s2);
9649566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(ll, &l));
9655b2fa520SLois Curfman McInnes     for (i = 0; i < m; i++) {
9665b2fa520SLois Curfman McInnes       x = l[i];
967637a0070SStefano Zampini       v = vv + i;
9689371c9d4SSatish Balay       for (j = 0; j < n; j++) {
9699371c9d4SSatish Balay         (*v) *= x;
9709371c9d4SSatish Balay         v += lda;
9719371c9d4SSatish Balay       }
9725b2fa520SLois Curfman McInnes     }
9739566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(ll, &l));
9749566063dSJacob Faibussowitsch     PetscCall(PetscLogFlops(1.0 * n * m));
9755b2fa520SLois Curfman McInnes   }
9765b2fa520SLois Curfman McInnes   if (rr) {
977637a0070SStefano Zampini     const PetscScalar *ar;
978637a0070SStefano Zampini 
9799566063dSJacob Faibussowitsch     PetscCall(VecGetLocalSize(rr, &s3a));
98008401ef6SPierre Jolivet     PetscCheck(s3a == s3, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Right scaling vec non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT ".", s3a, s3);
9819566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(rr, &ar));
9826f08ad05SPierre Jolivet     if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
9839566063dSJacob Faibussowitsch     PetscCall(VecGetArray(mdn->lvec, &r));
9849566063dSJacob Faibussowitsch     PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, ar, r, MPI_REPLACE));
9859566063dSJacob Faibussowitsch     PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ar, r, MPI_REPLACE));
9869566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(rr, &ar));
9875b2fa520SLois Curfman McInnes     for (i = 0; i < n; i++) {
9885b2fa520SLois Curfman McInnes       x = r[i];
989637a0070SStefano Zampini       v = vv + i * lda;
9902205254eSKarl Rupp       for (j = 0; j < m; j++) (*v++) *= x;
9915b2fa520SLois Curfman McInnes     }
9929566063dSJacob Faibussowitsch     PetscCall(VecRestoreArray(mdn->lvec, &r));
9939566063dSJacob Faibussowitsch     PetscCall(PetscLogFlops(1.0 * n * m));
9945b2fa520SLois Curfman McInnes   }
9959566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(mdn->A, &vv));
9963ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9975b2fa520SLois Curfman McInnes }
9985b2fa520SLois Curfman McInnes 
99966976f2fSJacob Faibussowitsch static PetscErrorCode MatNorm_MPIDense(Mat A, NormType type, PetscReal *nrm)
1000d71ae5a4SJacob Faibussowitsch {
10013501a2bdSLois Curfman McInnes   Mat_MPIDense      *mdn = (Mat_MPIDense *)A->data;
100213f74950SBarry Smith   PetscInt           i, j;
10036497c311SBarry Smith   PetscMPIInt        size, iN;
1004329f5518SBarry Smith   PetscReal          sum = 0.0;
1005637a0070SStefano Zampini   const PetscScalar *av, *v;
10063501a2bdSLois Curfman McInnes 
10073a40ed3dSBarry Smith   PetscFunctionBegin;
10089566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(mdn->A, &av));
1009637a0070SStefano Zampini   v = av;
10109566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
1011616b8fbbSStefano Zampini   if (size == 1) {
10129566063dSJacob Faibussowitsch     PetscCall(MatNorm(mdn->A, type, nrm));
10133501a2bdSLois Curfman McInnes   } else {
10143501a2bdSLois Curfman McInnes     if (type == NORM_FROBENIUS) {
1015d0f46423SBarry Smith       for (i = 0; i < mdn->A->cmap->n * mdn->A->rmap->n; i++) {
10169371c9d4SSatish Balay         sum += PetscRealPart(PetscConj(*v) * (*v));
10179371c9d4SSatish Balay         v++;
10183501a2bdSLois Curfman McInnes       }
1019462c564dSBarry Smith       PetscCallMPI(MPIU_Allreduce(&sum, nrm, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
10208f1a2a5eSBarry Smith       *nrm = PetscSqrtReal(*nrm);
10219566063dSJacob Faibussowitsch       PetscCall(PetscLogFlops(2.0 * mdn->A->cmap->n * mdn->A->rmap->n));
10223a40ed3dSBarry Smith     } else if (type == NORM_1) {
1023329f5518SBarry Smith       PetscReal *tmp, *tmp2;
10249566063dSJacob Faibussowitsch       PetscCall(PetscCalloc2(A->cmap->N, &tmp, A->cmap->N, &tmp2));
1025064f8208SBarry Smith       *nrm = 0.0;
1026637a0070SStefano Zampini       v    = av;
1027d0f46423SBarry Smith       for (j = 0; j < mdn->A->cmap->n; j++) {
1028d0f46423SBarry Smith         for (i = 0; i < mdn->A->rmap->n; i++) {
10299371c9d4SSatish Balay           tmp[j] += PetscAbsScalar(*v);
10309371c9d4SSatish Balay           v++;
10313501a2bdSLois Curfman McInnes         }
10323501a2bdSLois Curfman McInnes       }
10336497c311SBarry Smith       PetscCall(PetscMPIIntCast(A->cmap->N, &iN));
1034462c564dSBarry Smith       PetscCallMPI(MPIU_Allreduce(tmp, tmp2, iN, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
1035d0f46423SBarry Smith       for (j = 0; j < A->cmap->N; j++) {
1036064f8208SBarry Smith         if (tmp2[j] > *nrm) *nrm = tmp2[j];
10373501a2bdSLois Curfman McInnes       }
10389566063dSJacob Faibussowitsch       PetscCall(PetscFree2(tmp, tmp2));
10399566063dSJacob Faibussowitsch       PetscCall(PetscLogFlops(A->cmap->n * A->rmap->n));
10403a40ed3dSBarry Smith     } else if (type == NORM_INFINITY) { /* max row norm */
1041329f5518SBarry Smith       PetscReal ntemp;
10429566063dSJacob Faibussowitsch       PetscCall(MatNorm(mdn->A, type, &ntemp));
1043462c564dSBarry Smith       PetscCallMPI(MPIU_Allreduce(&ntemp, nrm, 1, MPIU_REAL, MPIU_MAX, PetscObjectComm((PetscObject)A)));
1044ce94432eSBarry Smith     } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "No support for two norm");
10453501a2bdSLois Curfman McInnes   }
10469566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(mdn->A, &av));
10473ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
10483501a2bdSLois Curfman McInnes }
10493501a2bdSLois Curfman McInnes 
105066976f2fSJacob Faibussowitsch static PetscErrorCode MatTranspose_MPIDense(Mat A, MatReuse reuse, Mat *matout)
1051d71ae5a4SJacob Faibussowitsch {
10523501a2bdSLois Curfman McInnes   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
10533501a2bdSLois Curfman McInnes   Mat           B;
1054d0f46423SBarry Smith   PetscInt      M = A->rmap->N, N = A->cmap->N, m, n, *rwork, rstart = A->rmap->rstart;
1055637a0070SStefano Zampini   PetscInt      j, i, lda;
105687828ca2SBarry Smith   PetscScalar  *v;
10573501a2bdSLois Curfman McInnes 
10583a40ed3dSBarry Smith   PetscFunctionBegin;
10597fb60732SBarry Smith   if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(A, *matout));
1060cf37664fSBarry Smith   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX) {
10619566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
10629566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(B, A->cmap->n, A->rmap->n, N, M));
10639566063dSJacob Faibussowitsch     PetscCall(MatSetType(B, ((PetscObject)A)->type_name));
10649566063dSJacob Faibussowitsch     PetscCall(MatMPIDenseSetPreallocation(B, NULL));
1065637a0070SStefano Zampini   } else B = *matout;
10663501a2bdSLois Curfman McInnes 
10679371c9d4SSatish Balay   m = a->A->rmap->n;
10689371c9d4SSatish Balay   n = a->A->cmap->n;
10699566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)&v));
10709566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
10719566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(m, &rwork));
10723501a2bdSLois Curfman McInnes   for (i = 0; i < m; i++) rwork[i] = rstart + i;
10731acff37aSSatish Balay   for (j = 0; j < n; j++) {
10749566063dSJacob Faibussowitsch     PetscCall(MatSetValues(B, 1, &j, m, rwork, v, INSERT_VALUES));
10758e3a54c0SPierre Jolivet     v = PetscSafePointerPlusOffset(v, lda);
10763501a2bdSLois Curfman McInnes   }
10779566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)&v));
10789566063dSJacob Faibussowitsch   PetscCall(PetscFree(rwork));
10799566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
10809566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
1081cf37664fSBarry Smith   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_REUSE_MATRIX) {
10823501a2bdSLois Curfman McInnes     *matout = B;
10833501a2bdSLois Curfman McInnes   } else {
10849566063dSJacob Faibussowitsch     PetscCall(MatHeaderMerge(A, &B));
10853501a2bdSLois Curfman McInnes   }
10863ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1087096963f5SLois Curfman McInnes }
1088096963f5SLois Curfman McInnes 
10896849ba73SBarry Smith static PetscErrorCode       MatDuplicate_MPIDense(Mat, MatDuplicateOption, Mat *);
109052c5f739Sprj- PETSC_INTERN PetscErrorCode MatScale_MPIDense(Mat, PetscScalar);
10918965ea79SLois Curfman McInnes 
109266976f2fSJacob Faibussowitsch static PetscErrorCode MatSetUp_MPIDense(Mat A)
1093d71ae5a4SJacob Faibussowitsch {
1094273d9f13SBarry Smith   PetscFunctionBegin;
10959566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(A->rmap));
10969566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(A->cmap));
109748a46eb9SPierre Jolivet   if (!A->preallocated) PetscCall(MatMPIDenseSetPreallocation(A, NULL));
10983ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1099273d9f13SBarry Smith }
1100273d9f13SBarry Smith 
110166976f2fSJacob Faibussowitsch static PetscErrorCode MatAXPY_MPIDense(Mat Y, PetscScalar alpha, Mat X, MatStructure str)
1102d71ae5a4SJacob Faibussowitsch {
1103488007eeSBarry Smith   Mat_MPIDense *A = (Mat_MPIDense *)Y->data, *B = (Mat_MPIDense *)X->data;
1104488007eeSBarry Smith 
1105488007eeSBarry Smith   PetscFunctionBegin;
11069566063dSJacob Faibussowitsch   PetscCall(MatAXPY(A->A, alpha, B->A, str));
11073ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1108488007eeSBarry Smith }
1109488007eeSBarry Smith 
111066976f2fSJacob Faibussowitsch static PetscErrorCode MatConjugate_MPIDense(Mat mat)
1111d71ae5a4SJacob Faibussowitsch {
1112ba337c44SJed Brown   Mat_MPIDense *a = (Mat_MPIDense *)mat->data;
1113ba337c44SJed Brown 
1114ba337c44SJed Brown   PetscFunctionBegin;
11159566063dSJacob Faibussowitsch   PetscCall(MatConjugate(a->A));
11163ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1117ba337c44SJed Brown }
1118ba337c44SJed Brown 
111966976f2fSJacob Faibussowitsch static PetscErrorCode MatRealPart_MPIDense(Mat A)
1120d71ae5a4SJacob Faibussowitsch {
1121ba337c44SJed Brown   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1122ba337c44SJed Brown 
1123ba337c44SJed Brown   PetscFunctionBegin;
11249566063dSJacob Faibussowitsch   PetscCall(MatRealPart(a->A));
11253ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1126ba337c44SJed Brown }
1127ba337c44SJed Brown 
112866976f2fSJacob Faibussowitsch static PetscErrorCode MatImaginaryPart_MPIDense(Mat A)
1129d71ae5a4SJacob Faibussowitsch {
1130ba337c44SJed Brown   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1131ba337c44SJed Brown 
1132ba337c44SJed Brown   PetscFunctionBegin;
11339566063dSJacob Faibussowitsch   PetscCall(MatImaginaryPart(a->A));
11343ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1135ba337c44SJed Brown }
1136ba337c44SJed Brown 
1137d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatGetColumnVector_MPIDense(Mat A, Vec v, PetscInt col)
1138d71ae5a4SJacob Faibussowitsch {
1139637a0070SStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
114049a6ff4bSBarry Smith 
114149a6ff4bSBarry Smith   PetscFunctionBegin;
114228b400f6SJacob Faibussowitsch   PetscCheck(a->A, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing local matrix");
114328b400f6SJacob Faibussowitsch   PetscCheck(a->A->ops->getcolumnvector, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing get column operation");
11449566063dSJacob Faibussowitsch   PetscCall((*a->A->ops->getcolumnvector)(a->A, v, col));
11453ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
114649a6ff4bSBarry Smith }
114749a6ff4bSBarry Smith 
1148857cbf51SRichard Tran Mills PETSC_INTERN PetscErrorCode MatGetColumnReductions_SeqDense(Mat, PetscInt, PetscReal *);
114952c5f739Sprj- 
115066976f2fSJacob Faibussowitsch static PetscErrorCode MatGetColumnReductions_MPIDense(Mat A, PetscInt type, PetscReal *reductions)
1151d71ae5a4SJacob Faibussowitsch {
1152a873a8cdSSam Reynolds   PetscInt      i, m, n;
11530716a85fSBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
11540716a85fSBarry Smith   PetscReal    *work;
11556497c311SBarry Smith   PetscMPIInt   in;
11560716a85fSBarry Smith 
11570716a85fSBarry Smith   PetscFunctionBegin;
11589566063dSJacob Faibussowitsch   PetscCall(MatGetSize(A, &m, &n));
11599566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(n, &work));
1160857cbf51SRichard Tran Mills   if (type == REDUCTION_MEAN_REALPART) {
11619566063dSJacob Faibussowitsch     PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_REALPART, work));
1162857cbf51SRichard Tran Mills   } else if (type == REDUCTION_MEAN_IMAGINARYPART) {
11639566063dSJacob Faibussowitsch     PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_IMAGINARYPART, work));
1164a873a8cdSSam Reynolds   } else {
11659566063dSJacob Faibussowitsch     PetscCall(MatGetColumnReductions_SeqDense(a->A, type, work));
1166a873a8cdSSam Reynolds   }
1167857cbf51SRichard Tran Mills   if (type == NORM_2) {
11680716a85fSBarry Smith     for (i = 0; i < n; i++) work[i] *= work[i];
11690716a85fSBarry Smith   }
11706497c311SBarry Smith   PetscCall(PetscMPIIntCast(n, &in));
1171857cbf51SRichard Tran Mills   if (type == NORM_INFINITY) {
1172462c564dSBarry Smith     PetscCallMPI(MPIU_Allreduce(work, reductions, in, MPIU_REAL, MPIU_MAX, A->hdr.comm));
11730716a85fSBarry Smith   } else {
1174462c564dSBarry Smith     PetscCallMPI(MPIU_Allreduce(work, reductions, in, MPIU_REAL, MPIU_SUM, A->hdr.comm));
11750716a85fSBarry Smith   }
11769566063dSJacob Faibussowitsch   PetscCall(PetscFree(work));
1177857cbf51SRichard Tran Mills   if (type == NORM_2) {
1178a873a8cdSSam Reynolds     for (i = 0; i < n; i++) reductions[i] = PetscSqrtReal(reductions[i]);
1179857cbf51SRichard Tran Mills   } else if (type == REDUCTION_MEAN_REALPART || type == REDUCTION_MEAN_IMAGINARYPART) {
1180a873a8cdSSam Reynolds     for (i = 0; i < n; i++) reductions[i] /= m;
11810716a85fSBarry Smith   }
11823ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
11830716a85fSBarry Smith }
11840716a85fSBarry Smith 
1185d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatSetRandom_MPIDense(Mat x, PetscRandom rctx)
1186d71ae5a4SJacob Faibussowitsch {
118773a71a0fSBarry Smith   Mat_MPIDense *d = (Mat_MPIDense *)x->data;
118873a71a0fSBarry Smith 
118973a71a0fSBarry Smith   PetscFunctionBegin;
11909566063dSJacob Faibussowitsch   PetscCall(MatSetRandom(d->A, rctx));
11913faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE)
11923faff063SStefano Zampini   x->offloadmask = d->A->offloadmask;
11933faff063SStefano Zampini #endif
11943ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
119573a71a0fSBarry Smith }
119673a71a0fSBarry Smith 
1197d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMissingDiagonal_MPIDense(Mat A, PetscBool *missing, PetscInt *d)
1198d71ae5a4SJacob Faibussowitsch {
11993b49f96aSBarry Smith   PetscFunctionBegin;
12003b49f96aSBarry Smith   *missing = PETSC_FALSE;
12013ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
12023b49f96aSBarry Smith }
12033b49f96aSBarry Smith 
12044222ddf1SHong Zhang static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat);
1205cc48ffa7SToby Isaac static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat);
12066718818eSStefano Zampini static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat);
1207a1f9a5e6SStefano Zampini static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat);
12086718818eSStefano Zampini static PetscErrorCode MatEqual_MPIDense(Mat, Mat, PetscBool *);
12096718818eSStefano Zampini static PetscErrorCode MatLoad_MPIDense(Mat, PetscViewer);
12105d8c7819SPierre Jolivet static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat);
1211cc48ffa7SToby Isaac 
121209dc0095SBarry Smith static struct _MatOps MatOps_Values = {MatSetValues_MPIDense,
121309dc0095SBarry Smith                                        MatGetRow_MPIDense,
121409dc0095SBarry Smith                                        MatRestoreRow_MPIDense,
121509dc0095SBarry Smith                                        MatMult_MPIDense,
121697304618SKris Buschelman                                        /*  4*/ MatMultAdd_MPIDense,
12177c922b88SBarry Smith                                        MatMultTranspose_MPIDense,
12187c922b88SBarry Smith                                        MatMultTransposeAdd_MPIDense,
1219f4259b30SLisandro Dalcin                                        NULL,
1220f4259b30SLisandro Dalcin                                        NULL,
1221f4259b30SLisandro Dalcin                                        NULL,
1222f4259b30SLisandro Dalcin                                        /* 10*/ NULL,
1223f4259b30SLisandro Dalcin                                        NULL,
1224f4259b30SLisandro Dalcin                                        NULL,
1225f4259b30SLisandro Dalcin                                        NULL,
122609dc0095SBarry Smith                                        MatTranspose_MPIDense,
122797304618SKris Buschelman                                        /* 15*/ MatGetInfo_MPIDense,
12286e4ee0c6SHong Zhang                                        MatEqual_MPIDense,
122909dc0095SBarry Smith                                        MatGetDiagonal_MPIDense,
12305b2fa520SLois Curfman McInnes                                        MatDiagonalScale_MPIDense,
123109dc0095SBarry Smith                                        MatNorm_MPIDense,
123297304618SKris Buschelman                                        /* 20*/ MatAssemblyBegin_MPIDense,
123309dc0095SBarry Smith                                        MatAssemblyEnd_MPIDense,
123409dc0095SBarry Smith                                        MatSetOption_MPIDense,
123509dc0095SBarry Smith                                        MatZeroEntries_MPIDense,
1236d519adbfSMatthew Knepley                                        /* 24*/ MatZeroRows_MPIDense,
1237f4259b30SLisandro Dalcin                                        NULL,
1238f4259b30SLisandro Dalcin                                        NULL,
1239f4259b30SLisandro Dalcin                                        NULL,
1240f4259b30SLisandro Dalcin                                        NULL,
12414994cf47SJed Brown                                        /* 29*/ MatSetUp_MPIDense,
1242f4259b30SLisandro Dalcin                                        NULL,
1243f4259b30SLisandro Dalcin                                        NULL,
1244c56a70eeSBarry Smith                                        MatGetDiagonalBlock_MPIDense,
1245f4259b30SLisandro Dalcin                                        NULL,
1246d519adbfSMatthew Knepley                                        /* 34*/ MatDuplicate_MPIDense,
1247f4259b30SLisandro Dalcin                                        NULL,
1248f4259b30SLisandro Dalcin                                        NULL,
1249f4259b30SLisandro Dalcin                                        NULL,
1250f4259b30SLisandro Dalcin                                        NULL,
1251d519adbfSMatthew Knepley                                        /* 39*/ MatAXPY_MPIDense,
12527dae84e0SHong Zhang                                        MatCreateSubMatrices_MPIDense,
1253f4259b30SLisandro Dalcin                                        NULL,
125409dc0095SBarry Smith                                        MatGetValues_MPIDense,
1255a76f77c3SStefano Zampini                                        MatCopy_MPIDense,
1256f4259b30SLisandro Dalcin                                        /* 44*/ NULL,
125709dc0095SBarry Smith                                        MatScale_MPIDense,
12582f605a99SJose E. Roman                                        MatShift_MPIDense,
1259f4259b30SLisandro Dalcin                                        NULL,
1260f4259b30SLisandro Dalcin                                        NULL,
126173a71a0fSBarry Smith                                        /* 49*/ MatSetRandom_MPIDense,
1262f4259b30SLisandro Dalcin                                        NULL,
1263f4259b30SLisandro Dalcin                                        NULL,
1264f4259b30SLisandro Dalcin                                        NULL,
1265f4259b30SLisandro Dalcin                                        NULL,
1266f4259b30SLisandro Dalcin                                        /* 54*/ NULL,
1267f4259b30SLisandro Dalcin                                        NULL,
1268f4259b30SLisandro Dalcin                                        NULL,
1269f4259b30SLisandro Dalcin                                        NULL,
1270f4259b30SLisandro Dalcin                                        NULL,
12717dae84e0SHong Zhang                                        /* 59*/ MatCreateSubMatrix_MPIDense,
1272b9b97703SBarry Smith                                        MatDestroy_MPIDense,
1273b9b97703SBarry Smith                                        MatView_MPIDense,
1274f4259b30SLisandro Dalcin                                        NULL,
1275f4259b30SLisandro Dalcin                                        NULL,
1276f4259b30SLisandro Dalcin                                        /* 64*/ NULL,
1277f4259b30SLisandro Dalcin                                        NULL,
1278f4259b30SLisandro Dalcin                                        NULL,
1279f4259b30SLisandro Dalcin                                        NULL,
1280f4259b30SLisandro Dalcin                                        NULL,
1281f4259b30SLisandro Dalcin                                        /* 69*/ NULL,
1282f4259b30SLisandro Dalcin                                        NULL,
1283f4259b30SLisandro Dalcin                                        NULL,
1284f4259b30SLisandro Dalcin                                        NULL,
1285f4259b30SLisandro Dalcin                                        NULL,
1286f4259b30SLisandro Dalcin                                        /* 74*/ NULL,
1287f4259b30SLisandro Dalcin                                        NULL,
1288f4259b30SLisandro Dalcin                                        NULL,
1289f4259b30SLisandro Dalcin                                        NULL,
1290f4259b30SLisandro Dalcin                                        NULL,
1291f4259b30SLisandro Dalcin                                        /* 79*/ NULL,
1292f4259b30SLisandro Dalcin                                        NULL,
1293f4259b30SLisandro Dalcin                                        NULL,
1294f4259b30SLisandro Dalcin                                        NULL,
12955bba2384SShri Abhyankar                                        /* 83*/ MatLoad_MPIDense,
1296f4259b30SLisandro Dalcin                                        NULL,
1297f4259b30SLisandro Dalcin                                        NULL,
1298f4259b30SLisandro Dalcin                                        NULL,
1299f4259b30SLisandro Dalcin                                        NULL,
1300f4259b30SLisandro Dalcin                                        NULL,
1301f4259b30SLisandro Dalcin                                        /* 89*/ NULL,
1302f4259b30SLisandro Dalcin                                        NULL,
1303f4259b30SLisandro Dalcin                                        NULL,
1304f4259b30SLisandro Dalcin                                        NULL,
1305f4259b30SLisandro Dalcin                                        NULL,
1306f4259b30SLisandro Dalcin                                        /* 94*/ NULL,
1307f4259b30SLisandro Dalcin                                        NULL,
13086718818eSStefano Zampini                                        MatMatTransposeMultSymbolic_MPIDense_MPIDense,
1309cc48ffa7SToby Isaac                                        MatMatTransposeMultNumeric_MPIDense_MPIDense,
1310f4259b30SLisandro Dalcin                                        NULL,
13114222ddf1SHong Zhang                                        /* 99*/ MatProductSetFromOptions_MPIDense,
1312f4259b30SLisandro Dalcin                                        NULL,
1313f4259b30SLisandro Dalcin                                        NULL,
1314ba337c44SJed Brown                                        MatConjugate_MPIDense,
1315f4259b30SLisandro Dalcin                                        NULL,
1316f4259b30SLisandro Dalcin                                        /*104*/ NULL,
1317ba337c44SJed Brown                                        MatRealPart_MPIDense,
1318ba337c44SJed Brown                                        MatImaginaryPart_MPIDense,
1319f4259b30SLisandro Dalcin                                        NULL,
1320f4259b30SLisandro Dalcin                                        NULL,
1321f4259b30SLisandro Dalcin                                        /*109*/ NULL,
1322f4259b30SLisandro Dalcin                                        NULL,
1323f4259b30SLisandro Dalcin                                        NULL,
132449a6ff4bSBarry Smith                                        MatGetColumnVector_MPIDense,
13253b49f96aSBarry Smith                                        MatMissingDiagonal_MPIDense,
1326f4259b30SLisandro Dalcin                                        /*114*/ NULL,
1327f4259b30SLisandro Dalcin                                        NULL,
1328f4259b30SLisandro Dalcin                                        NULL,
1329f4259b30SLisandro Dalcin                                        NULL,
1330f4259b30SLisandro Dalcin                                        NULL,
1331f4259b30SLisandro Dalcin                                        /*119*/ NULL,
1332f4259b30SLisandro Dalcin                                        NULL,
1333459e8d23SBlanca Mellado Pinto                                        MatMultHermitianTranspose_MPIDense,
1334459e8d23SBlanca Mellado Pinto                                        MatMultHermitianTransposeAdd_MPIDense,
1335f4259b30SLisandro Dalcin                                        NULL,
1336f4259b30SLisandro Dalcin                                        /*124*/ NULL,
1337a873a8cdSSam Reynolds                                        MatGetColumnReductions_MPIDense,
1338f4259b30SLisandro Dalcin                                        NULL,
1339f4259b30SLisandro Dalcin                                        NULL,
1340f4259b30SLisandro Dalcin                                        NULL,
1341f4259b30SLisandro Dalcin                                        /*129*/ NULL,
1342f4259b30SLisandro Dalcin                                        NULL,
13436718818eSStefano Zampini                                        MatTransposeMatMultSymbolic_MPIDense_MPIDense,
1344cb20be35SHong Zhang                                        MatTransposeMatMultNumeric_MPIDense_MPIDense,
1345f4259b30SLisandro Dalcin                                        NULL,
1346f4259b30SLisandro Dalcin                                        /*134*/ NULL,
1347f4259b30SLisandro Dalcin                                        NULL,
1348f4259b30SLisandro Dalcin                                        NULL,
1349f4259b30SLisandro Dalcin                                        NULL,
1350f4259b30SLisandro Dalcin                                        NULL,
1351f4259b30SLisandro Dalcin                                        /*139*/ NULL,
1352f4259b30SLisandro Dalcin                                        NULL,
1353f4259b30SLisandro Dalcin                                        NULL,
1354f4259b30SLisandro Dalcin                                        NULL,
1355f4259b30SLisandro Dalcin                                        NULL,
13564222ddf1SHong Zhang                                        MatCreateMPIMatConcatenateSeqMat_MPIDense,
1357f4259b30SLisandro Dalcin                                        /*145*/ NULL,
1358f4259b30SLisandro Dalcin                                        NULL,
135999a7f59eSMark Adams                                        NULL,
136099a7f59eSMark Adams                                        NULL,
13617fb60732SBarry Smith                                        NULL,
1362dec0b466SHong Zhang                                        /*150*/ NULL,
1363eede4a3fSMark Adams                                        NULL,
13644cc2b5b5SPierre Jolivet                                        NULL,
136542ce410bSJunchao Zhang                                        NULL,
136642ce410bSJunchao Zhang                                        NULL,
1367dec0b466SHong Zhang                                        NULL};
13688965ea79SLois Curfman McInnes 
136966976f2fSJacob Faibussowitsch static PetscErrorCode MatMPIDenseSetPreallocation_MPIDense(Mat mat, PetscScalar *data)
1370d71ae5a4SJacob Faibussowitsch {
1371637a0070SStefano Zampini   Mat_MPIDense *a     = (Mat_MPIDense *)mat->data;
137247d993e7Ssuyashtn   MatType       mtype = MATSEQDENSE;
1373a23d5eceSKris Buschelman 
1374a23d5eceSKris Buschelman   PetscFunctionBegin;
137528b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
13769566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(mat->rmap));
13779566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(mat->cmap));
1378637a0070SStefano Zampini   if (!a->A) {
13799566063dSJacob Faibussowitsch     PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
13809566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(a->A, mat->rmap->n, mat->cmap->N, mat->rmap->n, mat->cmap->N));
1381637a0070SStefano Zampini   }
138247d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA)
138347d993e7Ssuyashtn   PetscBool iscuda;
13849566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSECUDA, &iscuda));
138547d993e7Ssuyashtn   if (iscuda) mtype = MATSEQDENSECUDA;
138647d993e7Ssuyashtn #endif
138747d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
138847d993e7Ssuyashtn   PetscBool iship;
138947d993e7Ssuyashtn   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSEHIP, &iship));
139047d993e7Ssuyashtn   if (iship) mtype = MATSEQDENSEHIP;
139147d993e7Ssuyashtn #endif
139247d993e7Ssuyashtn   PetscCall(MatSetType(a->A, mtype));
13939566063dSJacob Faibussowitsch   PetscCall(MatSeqDenseSetPreallocation(a->A, data));
139447d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
13953faff063SStefano Zampini   mat->offloadmask = a->A->offloadmask;
13963faff063SStefano Zampini #endif
1397637a0070SStefano Zampini   mat->preallocated = PETSC_TRUE;
13986f08ad05SPierre Jolivet   mat->assembled    = PETSC_TRUE;
13993ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1400a23d5eceSKris Buschelman }
1401a23d5eceSKris Buschelman 
1402d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode MatConvert_MPIAIJ_MPIDense(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1403d71ae5a4SJacob Faibussowitsch {
14044e29119aSPierre Jolivet   Mat B, C;
14054e29119aSPierre Jolivet 
14064e29119aSPierre Jolivet   PetscFunctionBegin;
14079566063dSJacob Faibussowitsch   PetscCall(MatMPIAIJGetLocalMat(A, MAT_INITIAL_MATRIX, &C));
14089566063dSJacob Faibussowitsch   PetscCall(MatConvert_SeqAIJ_SeqDense(C, MATSEQDENSE, MAT_INITIAL_MATRIX, &B));
14099566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&C));
14104e29119aSPierre Jolivet   if (reuse == MAT_REUSE_MATRIX) {
14114e29119aSPierre Jolivet     C = *newmat;
14124e29119aSPierre Jolivet   } else C = NULL;
14139566063dSJacob Faibussowitsch   PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C));
14149566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&B));
14154e29119aSPierre Jolivet   if (reuse == MAT_INPLACE_MATRIX) {
14169566063dSJacob Faibussowitsch     PetscCall(MatHeaderReplace(A, &C));
14174e29119aSPierre Jolivet   } else if (reuse == MAT_INITIAL_MATRIX) *newmat = C;
14183ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14194e29119aSPierre Jolivet }
14204e29119aSPierre Jolivet 
142166976f2fSJacob Faibussowitsch static PetscErrorCode MatConvert_MPIDense_MPIAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1422d71ae5a4SJacob Faibussowitsch {
14234e29119aSPierre Jolivet   Mat B, C;
14244e29119aSPierre Jolivet 
14254e29119aSPierre Jolivet   PetscFunctionBegin;
14269566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLocalMatrix(A, &C));
14279566063dSJacob Faibussowitsch   PetscCall(MatConvert_SeqDense_SeqAIJ(C, MATSEQAIJ, MAT_INITIAL_MATRIX, &B));
14284e29119aSPierre Jolivet   if (reuse == MAT_REUSE_MATRIX) {
14294e29119aSPierre Jolivet     C = *newmat;
14304e29119aSPierre Jolivet   } else C = NULL;
14319566063dSJacob Faibussowitsch   PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C));
14329566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&B));
14334e29119aSPierre Jolivet   if (reuse == MAT_INPLACE_MATRIX) {
14349566063dSJacob Faibussowitsch     PetscCall(MatHeaderReplace(A, &C));
14354e29119aSPierre Jolivet   } else if (reuse == MAT_INITIAL_MATRIX) *newmat = C;
14363ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14374e29119aSPierre Jolivet }
14384e29119aSPierre Jolivet 
143965b80a83SHong Zhang #if defined(PETSC_HAVE_ELEMENTAL)
1440d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode MatConvert_MPIDense_Elemental(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1441d71ae5a4SJacob Faibussowitsch {
1442f709e8d7SPierre Jolivet   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
14438ea901baSHong Zhang   Mat           mat_elemental;
144432d7a744SHong Zhang   PetscScalar  *v;
1445f709e8d7SPierre Jolivet   PetscInt      m = A->rmap->n, N = A->cmap->N, rstart = A->rmap->rstart, i, *rows, *cols, lda;
14468ea901baSHong Zhang 
14478baccfbdSHong Zhang   PetscFunctionBegin;
1448378336b6SHong Zhang   if (reuse == MAT_REUSE_MATRIX) {
1449378336b6SHong Zhang     mat_elemental = *newmat;
14509566063dSJacob Faibussowitsch     PetscCall(MatZeroEntries(*newmat));
1451378336b6SHong Zhang   } else {
14529566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat_elemental));
14539566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(mat_elemental, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N));
14549566063dSJacob Faibussowitsch     PetscCall(MatSetType(mat_elemental, MATELEMENTAL));
14559566063dSJacob Faibussowitsch     PetscCall(MatSetUp(mat_elemental));
14569566063dSJacob Faibussowitsch     PetscCall(MatSetOption(mat_elemental, MAT_ROW_ORIENTED, PETSC_FALSE));
1457378336b6SHong Zhang   }
1458378336b6SHong Zhang 
14599566063dSJacob Faibussowitsch   PetscCall(PetscMalloc2(m, &rows, N, &cols));
146032d7a744SHong Zhang   for (i = 0; i < N; i++) cols[i] = i;
146132d7a744SHong Zhang   for (i = 0; i < m; i++) rows[i] = rstart + i;
14628ea901baSHong Zhang 
1463637a0070SStefano Zampini   /* PETSc-Elemental interface uses axpy for setting off-processor entries, only ADD_VALUES is allowed */
14649566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(A, &v));
1465f709e8d7SPierre Jolivet   PetscCall(MatDenseGetLDA(a->A, &lda));
1466f709e8d7SPierre Jolivet   if (lda == m) PetscCall(MatSetValues(mat_elemental, m, rows, N, cols, v, ADD_VALUES));
1467f709e8d7SPierre Jolivet   else {
1468f709e8d7SPierre Jolivet     for (i = 0; i < N; i++) PetscCall(MatSetValues(mat_elemental, m, rows, 1, &i, v + lda * i, ADD_VALUES));
1469f709e8d7SPierre Jolivet   }
14709566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(mat_elemental, MAT_FINAL_ASSEMBLY));
14719566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(mat_elemental, MAT_FINAL_ASSEMBLY));
14729566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(A, &v));
14739566063dSJacob Faibussowitsch   PetscCall(PetscFree2(rows, cols));
14748ea901baSHong Zhang 
1475511c6705SHong Zhang   if (reuse == MAT_INPLACE_MATRIX) {
14769566063dSJacob Faibussowitsch     PetscCall(MatHeaderReplace(A, &mat_elemental));
14778ea901baSHong Zhang   } else {
14788ea901baSHong Zhang     *newmat = mat_elemental;
14798ea901baSHong Zhang   }
14803ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14818baccfbdSHong Zhang }
148265b80a83SHong Zhang #endif
14838baccfbdSHong Zhang 
1484d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetColumn_MPIDense(Mat A, PetscInt col, PetscScalar **vals)
1485d71ae5a4SJacob Faibussowitsch {
148686aefd0dSHong Zhang   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
148786aefd0dSHong Zhang 
148886aefd0dSHong Zhang   PetscFunctionBegin;
14899566063dSJacob Faibussowitsch   PetscCall(MatDenseGetColumn(mat->A, col, vals));
14903ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
149186aefd0dSHong Zhang }
149286aefd0dSHong Zhang 
1493d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseRestoreColumn_MPIDense(Mat A, PetscScalar **vals)
1494d71ae5a4SJacob Faibussowitsch {
149586aefd0dSHong Zhang   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
149686aefd0dSHong Zhang 
149786aefd0dSHong Zhang   PetscFunctionBegin;
14989566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreColumn(mat->A, vals));
14993ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
150086aefd0dSHong Zhang }
150186aefd0dSHong Zhang 
1502d71ae5a4SJacob Faibussowitsch PetscErrorCode MatCreateMPIMatConcatenateSeqMat_MPIDense(MPI_Comm comm, Mat inmat, PetscInt n, MatReuse scall, Mat *outmat)
1503d71ae5a4SJacob Faibussowitsch {
150494e2cb23SJakub Kruzik   Mat_MPIDense *mat;
150594e2cb23SJakub Kruzik   PetscInt      m, nloc, N;
150694e2cb23SJakub Kruzik 
150794e2cb23SJakub Kruzik   PetscFunctionBegin;
15089566063dSJacob Faibussowitsch   PetscCall(MatGetSize(inmat, &m, &N));
15099566063dSJacob Faibussowitsch   PetscCall(MatGetLocalSize(inmat, NULL, &nloc));
151094e2cb23SJakub Kruzik   if (scall == MAT_INITIAL_MATRIX) { /* symbolic phase */
151194e2cb23SJakub Kruzik     PetscInt sum;
151294e2cb23SJakub Kruzik 
151348a46eb9SPierre Jolivet     if (n == PETSC_DECIDE) PetscCall(PetscSplitOwnership(comm, &n, &N));
151494e2cb23SJakub Kruzik     /* Check sum(n) = N */
1515462c564dSBarry Smith     PetscCallMPI(MPIU_Allreduce(&n, &sum, 1, MPIU_INT, MPI_SUM, comm));
151608401ef6SPierre Jolivet     PetscCheck(sum == N, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Sum of local columns %" PetscInt_FMT " != global columns %" PetscInt_FMT, sum, N);
151794e2cb23SJakub Kruzik 
15189566063dSJacob Faibussowitsch     PetscCall(MatCreateDense(comm, m, n, PETSC_DETERMINE, N, NULL, outmat));
15199566063dSJacob Faibussowitsch     PetscCall(MatSetOption(*outmat, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
152094e2cb23SJakub Kruzik   }
152194e2cb23SJakub Kruzik 
152294e2cb23SJakub Kruzik   /* numeric phase */
152394e2cb23SJakub Kruzik   mat = (Mat_MPIDense *)(*outmat)->data;
15249566063dSJacob Faibussowitsch   PetscCall(MatCopy(inmat, mat->A, SAME_NONZERO_PATTERN));
15253ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
152694e2cb23SJakub Kruzik }
152794e2cb23SJakub Kruzik 
1528d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVec_MPIDense(Mat A, PetscInt col, Vec *v)
1529d71ae5a4SJacob Faibussowitsch {
15306947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15316947451fSStefano Zampini   PetscInt      lda;
15326947451fSStefano Zampini 
15336947451fSStefano Zampini   PetscFunctionBegin;
153428b400f6SJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
153528b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1536d16ceb75SStefano Zampini   if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
15376947451fSStefano Zampini   a->vecinuse = col + 1;
15389566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
15399566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(a->A, (PetscScalar **)&a->ptrinuse));
15408e3a54c0SPierre Jolivet   PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
15416947451fSStefano Zampini   *v = a->cvec;
15423ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15436947451fSStefano Zampini }
15446947451fSStefano Zampini 
1545d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVec_MPIDense(Mat A, PetscInt col, Vec *v)
1546d71ae5a4SJacob Faibussowitsch {
15476947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15486947451fSStefano Zampini 
15496947451fSStefano Zampini   PetscFunctionBegin;
155028b400f6SJacob Faibussowitsch   PetscCheck(a->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
155128b400f6SJacob Faibussowitsch   PetscCheck(a->cvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column vector");
15524186a4bbSPierre Jolivet   VecCheckAssembled(a->cvec);
15536947451fSStefano Zampini   a->vecinuse = 0;
15549566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(a->A, (PetscScalar **)&a->ptrinuse));
15559566063dSJacob Faibussowitsch   PetscCall(VecResetArray(a->cvec));
1556742765d3SMatthew Knepley   if (v) *v = NULL;
15573ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15586947451fSStefano Zampini }
15596947451fSStefano Zampini 
1560d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v)
1561d71ae5a4SJacob Faibussowitsch {
15626947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15636947451fSStefano Zampini   PetscInt      lda;
15646947451fSStefano Zampini 
15656947451fSStefano Zampini   PetscFunctionBegin;
156628b400f6SJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
156728b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1568d16ceb75SStefano Zampini   if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
15696947451fSStefano Zampini   a->vecinuse = col + 1;
15709566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
15719566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, &a->ptrinuse));
15728e3a54c0SPierre Jolivet   PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
15739566063dSJacob Faibussowitsch   PetscCall(VecLockReadPush(a->cvec));
15746947451fSStefano Zampini   *v = a->cvec;
15753ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15766947451fSStefano Zampini }
15776947451fSStefano Zampini 
1578d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v)
1579d71ae5a4SJacob Faibussowitsch {
15806947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15816947451fSStefano Zampini 
15826947451fSStefano Zampini   PetscFunctionBegin;
15835f80ce2aSJacob Faibussowitsch   PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
15845f80ce2aSJacob Faibussowitsch   PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector");
15854186a4bbSPierre Jolivet   VecCheckAssembled(a->cvec);
15866947451fSStefano Zampini   a->vecinuse = 0;
15879566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, &a->ptrinuse));
15889566063dSJacob Faibussowitsch   PetscCall(VecLockReadPop(a->cvec));
15899566063dSJacob Faibussowitsch   PetscCall(VecResetArray(a->cvec));
1590742765d3SMatthew Knepley   if (v) *v = NULL;
15913ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15926947451fSStefano Zampini }
15936947451fSStefano Zampini 
1594d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v)
1595d71ae5a4SJacob Faibussowitsch {
15966947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15976947451fSStefano Zampini   PetscInt      lda;
15986947451fSStefano Zampini 
15996947451fSStefano Zampini   PetscFunctionBegin;
16005f80ce2aSJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
16015f80ce2aSJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1602d16ceb75SStefano Zampini   if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
16036947451fSStefano Zampini   a->vecinuse = col + 1;
16049566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
16059566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(a->A, (PetscScalar **)&a->ptrinuse));
16068e3a54c0SPierre Jolivet   PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
16076947451fSStefano Zampini   *v = a->cvec;
16083ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
16096947451fSStefano Zampini }
16106947451fSStefano Zampini 
1611d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v)
1612d71ae5a4SJacob Faibussowitsch {
16136947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
16146947451fSStefano Zampini 
16156947451fSStefano Zampini   PetscFunctionBegin;
16165f80ce2aSJacob Faibussowitsch   PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
16175f80ce2aSJacob Faibussowitsch   PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector");
16184186a4bbSPierre Jolivet   VecCheckAssembled(a->cvec);
16196947451fSStefano Zampini   a->vecinuse = 0;
16209566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(a->A, (PetscScalar **)&a->ptrinuse));
16219566063dSJacob Faibussowitsch   PetscCall(VecResetArray(a->cvec));
1622742765d3SMatthew Knepley   if (v) *v = NULL;
16233ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
16246947451fSStefano Zampini }
16256947451fSStefano Zampini 
162666976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseGetSubMatrix_MPIDense(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *v)
1627d71ae5a4SJacob Faibussowitsch {
16285ea7661aSPierre Jolivet   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1629616b8fbbSStefano Zampini   Mat_MPIDense *c;
1630616b8fbbSStefano Zampini   MPI_Comm      comm;
1631a2748737SPierre Jolivet   PetscInt      pbegin, pend;
16325ea7661aSPierre Jolivet 
16335ea7661aSPierre Jolivet   PetscFunctionBegin;
16349566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
16355f80ce2aSJacob Faibussowitsch   PetscCheck(!a->vecinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
16365f80ce2aSJacob Faibussowitsch   PetscCheck(!a->matinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1637a2748737SPierre Jolivet   pbegin = PetscMax(0, PetscMin(A->rmap->rend, rbegin) - A->rmap->rstart);
1638a2748737SPierre Jolivet   pend   = PetscMin(A->rmap->n, PetscMax(0, rend - A->rmap->rstart));
16395ea7661aSPierre Jolivet   if (!a->cmat) {
16409566063dSJacob Faibussowitsch     PetscCall(MatCreate(comm, &a->cmat));
16419566063dSJacob Faibussowitsch     PetscCall(MatSetType(a->cmat, ((PetscObject)A)->type_name));
1642a2748737SPierre Jolivet     if (rend - rbegin == A->rmap->N) PetscCall(PetscLayoutReference(A->rmap, &a->cmat->rmap));
1643a2748737SPierre Jolivet     else {
1644a2748737SPierre Jolivet       PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, pend - pbegin));
1645a2748737SPierre Jolivet       PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin));
1646a2748737SPierre Jolivet       PetscCall(PetscLayoutSetUp(a->cmat->rmap));
1647a2748737SPierre Jolivet     }
16489566063dSJacob Faibussowitsch     PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin));
16499566063dSJacob Faibussowitsch     PetscCall(PetscLayoutSetUp(a->cmat->cmap));
1650a2748737SPierre Jolivet   } else {
1651a2748737SPierre Jolivet     PetscBool same = (PetscBool)(rend - rbegin == a->cmat->rmap->N);
1652a2748737SPierre Jolivet     if (same && a->cmat->rmap->N != A->rmap->N) {
1653a2748737SPierre Jolivet       same = (PetscBool)(pend - pbegin == a->cmat->rmap->n);
1654462c564dSBarry Smith       PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &same, 1, MPIU_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
1655a2748737SPierre Jolivet     }
1656a2748737SPierre Jolivet     if (!same) {
1657a2748737SPierre Jolivet       PetscCall(PetscLayoutDestroy(&a->cmat->rmap));
1658a2748737SPierre Jolivet       PetscCall(PetscLayoutCreate(comm, &a->cmat->rmap));
1659a2748737SPierre Jolivet       PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, pend - pbegin));
1660a2748737SPierre Jolivet       PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin));
1661a2748737SPierre Jolivet       PetscCall(PetscLayoutSetUp(a->cmat->rmap));
1662a2748737SPierre Jolivet     }
1663a2748737SPierre Jolivet     if (cend - cbegin != a->cmat->cmap->N) {
16649566063dSJacob Faibussowitsch       PetscCall(PetscLayoutDestroy(&a->cmat->cmap));
16659566063dSJacob Faibussowitsch       PetscCall(PetscLayoutCreate(comm, &a->cmat->cmap));
16669566063dSJacob Faibussowitsch       PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin));
16679566063dSJacob Faibussowitsch       PetscCall(PetscLayoutSetUp(a->cmat->cmap));
16685ea7661aSPierre Jolivet     }
1669a2748737SPierre Jolivet   }
1670616b8fbbSStefano Zampini   c = (Mat_MPIDense *)a->cmat->data;
16715f80ce2aSJacob Faibussowitsch   PetscCheck(!c->A, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1672a2748737SPierre Jolivet   PetscCall(MatDenseGetSubMatrix(a->A, pbegin, pend, cbegin, cend, &c->A));
16733faff063SStefano Zampini 
1674616b8fbbSStefano Zampini   a->cmat->preallocated = PETSC_TRUE;
1675616b8fbbSStefano Zampini   a->cmat->assembled    = PETSC_TRUE;
16763faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE)
16773faff063SStefano Zampini   a->cmat->offloadmask = c->A->offloadmask;
16783faff063SStefano Zampini #endif
16795ea7661aSPierre Jolivet   a->matinuse = cbegin + 1;
16805ea7661aSPierre Jolivet   *v          = a->cmat;
16813ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
16825ea7661aSPierre Jolivet }
16835ea7661aSPierre Jolivet 
168466976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreSubMatrix_MPIDense(Mat A, Mat *v)
1685d71ae5a4SJacob Faibussowitsch {
16865ea7661aSPierre Jolivet   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1687616b8fbbSStefano Zampini   Mat_MPIDense *c;
16885ea7661aSPierre Jolivet 
16895ea7661aSPierre Jolivet   PetscFunctionBegin;
16905f80ce2aSJacob Faibussowitsch   PetscCheck(a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetSubMatrix() first");
16915f80ce2aSJacob Faibussowitsch   PetscCheck(a->cmat, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal matrix");
16925f80ce2aSJacob Faibussowitsch   PetscCheck(*v == a->cmat, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Not the matrix obtained from MatDenseGetSubMatrix()");
16935ea7661aSPierre Jolivet   a->matinuse = 0;
1694616b8fbbSStefano Zampini   c           = (Mat_MPIDense *)a->cmat->data;
16959566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreSubMatrix(a->A, &c->A));
1696742765d3SMatthew Knepley   if (v) *v = NULL;
16973faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE)
16983faff063SStefano Zampini   A->offloadmask = a->A->offloadmask;
16993faff063SStefano Zampini #endif
17003ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
17015ea7661aSPierre Jolivet }
17025ea7661aSPierre Jolivet 
1703002bc6daSRichard Tran Mills /*MC
1704002bc6daSRichard Tran Mills    MATMPIDENSE - MATMPIDENSE = "mpidense" - A matrix type to be used for distributed dense matrices.
1705002bc6daSRichard Tran Mills 
17062ef1f0ffSBarry Smith    Options Database Key:
170711a5261eSBarry Smith . -mat_type mpidense - sets the matrix type to `MATMPIDENSE` during a call to `MatSetFromOptions()`
1708002bc6daSRichard Tran Mills 
1709002bc6daSRichard Tran Mills   Level: beginner
1710002bc6daSRichard Tran Mills 
17111cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MatCreateDense()`, `MATSEQDENSE`, `MATDENSE`
1712002bc6daSRichard Tran Mills M*/
17134742e46bSJacob Faibussowitsch PetscErrorCode MatCreate_MPIDense(Mat mat)
1714d71ae5a4SJacob Faibussowitsch {
1715273d9f13SBarry Smith   Mat_MPIDense *a;
1716273d9f13SBarry Smith 
1717273d9f13SBarry Smith   PetscFunctionBegin;
17184dfa11a4SJacob Faibussowitsch   PetscCall(PetscNew(&a));
1719b0a32e0cSBarry Smith   mat->data   = (void *)a;
1720aea10558SJacob Faibussowitsch   mat->ops[0] = MatOps_Values;
1721273d9f13SBarry Smith 
1722273d9f13SBarry Smith   mat->insertmode = NOT_SET_VALUES;
1723273d9f13SBarry Smith 
1724273d9f13SBarry Smith   /* build cache for off array entries formed */
1725273d9f13SBarry Smith   a->donotstash = PETSC_FALSE;
17262205254eSKarl Rupp 
17279566063dSJacob Faibussowitsch   PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)mat), 1, &mat->stash));
1728273d9f13SBarry Smith 
1729273d9f13SBarry Smith   /* stuff used for matrix vector multiply */
1730f4259b30SLisandro Dalcin   a->lvec        = NULL;
1731f4259b30SLisandro Dalcin   a->Mvctx       = NULL;
1732273d9f13SBarry Smith   a->roworiented = PETSC_TRUE;
1733273d9f13SBarry Smith 
17349566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", MatDenseGetLDA_MPIDense));
17359566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", MatDenseSetLDA_MPIDense));
17369566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", MatDenseGetArray_MPIDense));
17379566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", MatDenseRestoreArray_MPIDense));
17389566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", MatDenseGetArrayRead_MPIDense));
17399566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", MatDenseRestoreArrayRead_MPIDense));
17409566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", MatDenseGetArrayWrite_MPIDense));
17419566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", MatDenseRestoreArrayWrite_MPIDense));
17429566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", MatDensePlaceArray_MPIDense));
17439566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", MatDenseResetArray_MPIDense));
17449566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", MatDenseReplaceArray_MPIDense));
17459566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_MPIDense));
17469566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_MPIDense));
17479566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_MPIDense));
17489566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_MPIDense));
17499566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_MPIDense));
17509566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_MPIDense));
17519566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_MPIDense));
17529566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_MPIDense));
17539566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", MatConvert_MPIAIJ_MPIDense));
17549566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", MatConvert_MPIDense_MPIAIJ));
17558baccfbdSHong Zhang #if defined(PETSC_HAVE_ELEMENTAL)
17569566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", MatConvert_MPIDense_Elemental));
17578baccfbdSHong Zhang #endif
1758d24d4204SJose E. Roman #if defined(PETSC_HAVE_SCALAPACK)
17599566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", MatConvert_Dense_ScaLAPACK));
1760d24d4204SJose E. Roman #endif
1761637a0070SStefano Zampini #if defined(PETSC_HAVE_CUDA)
17629566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", MatConvert_MPIDense_MPIDenseCUDA));
1763637a0070SStefano Zampini #endif
17649566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", MatMPIDenseSetPreallocation_MPIDense));
17659566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
17669566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
17676718818eSStefano Zampini #if defined(PETSC_HAVE_CUDA)
17689566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
17699566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
17706718818eSStefano Zampini #endif
177147d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
177247d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", MatConvert_MPIDense_MPIDenseHIP));
177347d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
177447d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
177547d993e7Ssuyashtn #endif
17769566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", MatDenseGetColumn_MPIDense));
17779566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", MatDenseRestoreColumn_MPIDense));
17780be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", MatMultAddColumnRange_MPIDense));
17790be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", MatMultHermitianTransposeColumnRange_MPIDense));
17800be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", MatMultHermitianTransposeAddColumnRange_MPIDense));
17819566063dSJacob Faibussowitsch   PetscCall(PetscObjectChangeTypeName((PetscObject)mat, MATMPIDENSE));
17823ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1783273d9f13SBarry Smith }
1784273d9f13SBarry Smith 
1785209238afSKris Buschelman /*MC
1786002d173eSKris Buschelman    MATDENSE - MATDENSE = "dense" - A matrix type to be used for dense matrices.
1787209238afSKris Buschelman 
178811a5261eSBarry Smith    This matrix type is identical to `MATSEQDENSE` when constructed with a single process communicator,
178911a5261eSBarry Smith    and `MATMPIDENSE` otherwise.
1790209238afSKris Buschelman 
17912ef1f0ffSBarry Smith    Options Database Key:
179211a5261eSBarry Smith . -mat_type dense - sets the matrix type to `MATDENSE` during a call to `MatSetFromOptions()`
1793209238afSKris Buschelman 
1794209238afSKris Buschelman   Level: beginner
1795209238afSKris Buschelman 
17961cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATSEQDENSE`, `MATMPIDENSE`, `MATDENSECUDA`, `MATDENSEHIP`
17976947451fSStefano Zampini M*/
17986947451fSStefano Zampini 
17995d83a8b1SBarry Smith /*@
1800273d9f13SBarry Smith   MatMPIDenseSetPreallocation - Sets the array used to store the matrix entries
1801273d9f13SBarry Smith 
18028f490ea3SStefano Zampini   Collective
1803273d9f13SBarry Smith 
1804273d9f13SBarry Smith   Input Parameters:
1805fe59aa6dSJacob Faibussowitsch + B    - the matrix
18062ef1f0ffSBarry Smith - data - optional location of matrix data.  Set to `NULL` for PETSc
1807273d9f13SBarry Smith          to control all matrix memory allocation.
1808273d9f13SBarry Smith 
18092ef1f0ffSBarry Smith   Level: intermediate
18102ef1f0ffSBarry Smith 
1811273d9f13SBarry Smith   Notes:
18122ef1f0ffSBarry Smith   The dense format is fully compatible with standard Fortran
1813273d9f13SBarry Smith   storage by columns.
1814273d9f13SBarry Smith 
1815273d9f13SBarry Smith   The data input variable is intended primarily for Fortran programmers
1816273d9f13SBarry Smith   who wish to allocate their own matrix memory space.  Most users should
18172ef1f0ffSBarry Smith   set `data` to `NULL`.
1818273d9f13SBarry Smith 
18191cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATMPIDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()`
1820273d9f13SBarry Smith @*/
1821d71ae5a4SJacob Faibussowitsch PetscErrorCode MatMPIDenseSetPreallocation(Mat B, PetscScalar *data)
1822d71ae5a4SJacob Faibussowitsch {
1823273d9f13SBarry Smith   PetscFunctionBegin;
1824d5ea218eSStefano Zampini   PetscValidHeaderSpecific(B, MAT_CLASSID, 1);
1825cac4c232SBarry Smith   PetscTryMethod(B, "MatMPIDenseSetPreallocation_C", (Mat, PetscScalar *), (B, data));
18263ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1827273d9f13SBarry Smith }
1828273d9f13SBarry Smith 
1829d3042a70SBarry Smith /*@
183011a5261eSBarry Smith   MatDensePlaceArray - Allows one to replace the array in a `MATDENSE` matrix with an
1831d3042a70SBarry Smith   array provided by the user. This is useful to avoid copying an array
1832d3042a70SBarry Smith   into a matrix
1833d3042a70SBarry Smith 
1834d3042a70SBarry Smith   Not Collective
1835d3042a70SBarry Smith 
1836d3042a70SBarry Smith   Input Parameters:
1837d3042a70SBarry Smith + mat   - the matrix
1838d3042a70SBarry Smith - array - the array in column major order
1839d3042a70SBarry Smith 
18402ef1f0ffSBarry Smith   Level: developer
18412ef1f0ffSBarry Smith 
184211a5261eSBarry Smith   Note:
1843*4d5500e8SJunchao Zhang   Adding `const` to `array` was an oversight, see notes in `VecPlaceArray()`.
1844*4d5500e8SJunchao Zhang 
184511a5261eSBarry 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
1846d3042a70SBarry Smith   freed when the matrix is destroyed.
1847d3042a70SBarry Smith 
18481cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDenseResetArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`,
184911a5261eSBarry Smith           `MatDenseReplaceArray()`
1850d3042a70SBarry Smith @*/
1851d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDensePlaceArray(Mat mat, const PetscScalar *array)
1852d71ae5a4SJacob Faibussowitsch {
1853d3042a70SBarry Smith   PetscFunctionBegin;
1854d5ea218eSStefano Zampini   PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
1855cac4c232SBarry Smith   PetscUseMethod(mat, "MatDensePlaceArray_C", (Mat, const PetscScalar *), (mat, array));
18569566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)mat));
185747d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
1858637a0070SStefano Zampini   mat->offloadmask = PETSC_OFFLOAD_CPU;
1859637a0070SStefano Zampini #endif
18603ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1861d3042a70SBarry Smith }
1862d3042a70SBarry Smith 
1863d3042a70SBarry Smith /*@
186411a5261eSBarry Smith   MatDenseResetArray - Resets the matrix array to that it previously had before the call to `MatDensePlaceArray()`
1865d3042a70SBarry Smith 
1866d3042a70SBarry Smith   Not Collective
1867d3042a70SBarry Smith 
18682fe279fdSBarry Smith   Input Parameter:
1869d3042a70SBarry Smith . mat - the matrix
1870d3042a70SBarry Smith 
18712ef1f0ffSBarry Smith   Level: developer
18722ef1f0ffSBarry Smith 
187311a5261eSBarry Smith   Note:
187411a5261eSBarry Smith   You can only call this after a call to `MatDensePlaceArray()`
1875d3042a70SBarry Smith 
18761cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDensePlaceArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`
1877d3042a70SBarry Smith @*/
1878d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseResetArray(Mat mat)
1879d71ae5a4SJacob Faibussowitsch {
1880d3042a70SBarry Smith   PetscFunctionBegin;
1881d5ea218eSStefano Zampini   PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
1882cac4c232SBarry Smith   PetscUseMethod(mat, "MatDenseResetArray_C", (Mat), (mat));
18839566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)mat));
18843ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1885d3042a70SBarry Smith }
1886d3042a70SBarry Smith 
1887d5ea218eSStefano Zampini /*@
1888d5ea218eSStefano Zampini   MatDenseReplaceArray - Allows one to replace the array in a dense matrix with an
1889d5ea218eSStefano Zampini   array provided by the user. This is useful to avoid copying an array
1890d5ea218eSStefano Zampini   into a matrix
1891d5ea218eSStefano Zampini 
1892d5ea218eSStefano Zampini   Not Collective
1893d5ea218eSStefano Zampini 
1894d5ea218eSStefano Zampini   Input Parameters:
1895d5ea218eSStefano Zampini + mat   - the matrix
1896d5ea218eSStefano Zampini - array - the array in column major order
1897d5ea218eSStefano Zampini 
18982ef1f0ffSBarry Smith   Level: developer
18992ef1f0ffSBarry Smith 
190011a5261eSBarry Smith   Note:
1901*4d5500e8SJunchao Zhang   Adding `const` to `array` was an oversight, see notes in `VecPlaceArray()`.
1902*4d5500e8SJunchao Zhang 
190311a5261eSBarry Smith   The memory passed in MUST be obtained with `PetscMalloc()` and CANNOT be
1904d5ea218eSStefano Zampini   freed by the user. It will be freed when the matrix is destroyed.
1905d5ea218eSStefano Zampini 
19061cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MatDensePlaceArray()`, `MatDenseGetArray()`, `VecReplaceArray()`
1907d5ea218eSStefano Zampini @*/
1908d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseReplaceArray(Mat mat, const PetscScalar *array)
1909d71ae5a4SJacob Faibussowitsch {
1910d5ea218eSStefano Zampini   PetscFunctionBegin;
1911d5ea218eSStefano Zampini   PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
1912cac4c232SBarry Smith   PetscUseMethod(mat, "MatDenseReplaceArray_C", (Mat, const PetscScalar *), (mat, array));
19139566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)mat));
191447d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
1915d5ea218eSStefano Zampini   mat->offloadmask = PETSC_OFFLOAD_CPU;
1916d5ea218eSStefano Zampini #endif
19173ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1918d5ea218eSStefano Zampini }
1919d5ea218eSStefano Zampini 
19205d83a8b1SBarry Smith /*@
192111a5261eSBarry Smith   MatCreateDense - Creates a matrix in `MATDENSE` format.
19228965ea79SLois Curfman McInnes 
1923d083f849SBarry Smith   Collective
1924db81eaa0SLois Curfman McInnes 
19258965ea79SLois Curfman McInnes   Input Parameters:
1926db81eaa0SLois Curfman McInnes + comm - MPI communicator
19272ef1f0ffSBarry Smith . m    - number of local rows (or `PETSC_DECIDE` to have calculated if `M` is given)
19282ef1f0ffSBarry Smith . n    - number of local columns (or `PETSC_DECIDE` to have calculated if `N` is given)
19292ef1f0ffSBarry Smith . M    - number of global rows (or `PETSC_DECIDE` to have calculated if `m` is given)
19302ef1f0ffSBarry Smith . N    - number of global columns (or `PETSC_DECIDE` to have calculated if `n` is given)
19312ef1f0ffSBarry Smith - data - optional location of matrix data.  Set data to `NULL` (`PETSC_NULL_SCALAR` for Fortran users) for PETSc
1932dfc5480cSLois Curfman McInnes    to control all matrix memory allocation.
19338965ea79SLois Curfman McInnes 
19348965ea79SLois Curfman McInnes   Output Parameter:
1935477f1c0bSLois Curfman McInnes . A - the matrix
19368965ea79SLois Curfman McInnes 
19372ef1f0ffSBarry Smith   Level: intermediate
19382ef1f0ffSBarry Smith 
1939b259b22eSLois Curfman McInnes   Notes:
19402ef1f0ffSBarry Smith   The dense format is fully compatible with standard Fortran
194139ddd567SLois Curfman McInnes   storage by columns.
194211a5261eSBarry Smith 
194311a5261eSBarry Smith   Although local portions of the matrix are stored in column-major
1944002bc6daSRichard Tran Mills   order, the matrix is partitioned across MPI ranks by row.
19458965ea79SLois Curfman McInnes 
194618f449edSLois Curfman McInnes   The data input variable is intended primarily for Fortran programmers
194718f449edSLois Curfman McInnes   who wish to allocate their own matrix memory space.  Most users should
19482ef1f0ffSBarry Smith   set `data` to `NULL` (`PETSC_NULL_SCALAR` for Fortran users).
194918f449edSLois Curfman McInnes 
19508965ea79SLois Curfman McInnes   The user MUST specify either the local or global matrix dimensions
19518965ea79SLois Curfman McInnes   (possibly both).
19528965ea79SLois Curfman McInnes 
19531cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()`
19548965ea79SLois Curfman McInnes @*/
1955d71ae5a4SJacob Faibussowitsch PetscErrorCode MatCreateDense(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt M, PetscInt N, PetscScalar *data, Mat *A)
1956d71ae5a4SJacob Faibussowitsch {
19573a40ed3dSBarry Smith   PetscFunctionBegin;
19589566063dSJacob Faibussowitsch   PetscCall(MatCreate(comm, A));
19599566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(*A, m, n, M, N));
19606f08ad05SPierre Jolivet   PetscCall(MatSetType(*A, MATDENSE));
19619566063dSJacob Faibussowitsch   PetscCall(MatSeqDenseSetPreallocation(*A, data));
19626f08ad05SPierre Jolivet   PetscCall(MatMPIDenseSetPreallocation(*A, data));
19633ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
19648965ea79SLois Curfman McInnes }
19658965ea79SLois Curfman McInnes 
1966d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDuplicate_MPIDense(Mat A, MatDuplicateOption cpvalues, Mat *newmat)
1967d71ae5a4SJacob Faibussowitsch {
19688965ea79SLois Curfman McInnes   Mat           mat;
19693501a2bdSLois Curfman McInnes   Mat_MPIDense *a, *oldmat = (Mat_MPIDense *)A->data;
19708965ea79SLois Curfman McInnes 
19713a40ed3dSBarry Smith   PetscFunctionBegin;
1972f4259b30SLisandro Dalcin   *newmat = NULL;
19739566063dSJacob Faibussowitsch   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat));
19749566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(mat, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N));
19759566063dSJacob Faibussowitsch   PetscCall(MatSetType(mat, ((PetscObject)A)->type_name));
1976834f8fabSBarry Smith   a = (Mat_MPIDense *)mat->data;
19775aa7edbeSHong Zhang 
1978d5f3da31SBarry Smith   mat->factortype   = A->factortype;
1979c456f294SBarry Smith   mat->assembled    = PETSC_TRUE;
1980273d9f13SBarry Smith   mat->preallocated = PETSC_TRUE;
19818965ea79SLois Curfman McInnes 
1982e0fa3b82SLois Curfman McInnes   mat->insertmode = NOT_SET_VALUES;
19833782ba37SSatish Balay   a->donotstash   = oldmat->donotstash;
1984e04c1aa4SHong Zhang 
19859566063dSJacob Faibussowitsch   PetscCall(PetscLayoutReference(A->rmap, &mat->rmap));
19869566063dSJacob Faibussowitsch   PetscCall(PetscLayoutReference(A->cmap, &mat->cmap));
19878965ea79SLois Curfman McInnes 
19889566063dSJacob Faibussowitsch   PetscCall(MatDuplicate(oldmat->A, cpvalues, &a->A));
198901b82886SBarry Smith 
19908965ea79SLois Curfman McInnes   *newmat = mat;
19913ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
19928965ea79SLois Curfman McInnes }
19938965ea79SLois Curfman McInnes 
199466976f2fSJacob Faibussowitsch static PetscErrorCode MatLoad_MPIDense(Mat newMat, PetscViewer viewer)
1995d71ae5a4SJacob Faibussowitsch {
199687d5ce66SSatish Balay   PetscBool isbinary;
199787d5ce66SSatish Balay #if defined(PETSC_HAVE_HDF5)
199887d5ce66SSatish Balay   PetscBool ishdf5;
199987d5ce66SSatish Balay #endif
2000eb91f321SVaclav Hapla 
2001eb91f321SVaclav Hapla   PetscFunctionBegin;
2002eb91f321SVaclav Hapla   PetscValidHeaderSpecific(newMat, MAT_CLASSID, 1);
2003eb91f321SVaclav Hapla   PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
2004eb91f321SVaclav Hapla   /* force binary viewer to load .info file if it has not yet done so */
20059566063dSJacob Faibussowitsch   PetscCall(PetscViewerSetUp(viewer));
20069566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
200787d5ce66SSatish Balay #if defined(PETSC_HAVE_HDF5)
20089566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERHDF5, &ishdf5));
200987d5ce66SSatish Balay #endif
2010eb91f321SVaclav Hapla   if (isbinary) {
20119566063dSJacob Faibussowitsch     PetscCall(MatLoad_Dense_Binary(newMat, viewer));
2012eb91f321SVaclav Hapla #if defined(PETSC_HAVE_HDF5)
201387d5ce66SSatish Balay   } else if (ishdf5) {
20149566063dSJacob Faibussowitsch     PetscCall(MatLoad_Dense_HDF5(newMat, viewer));
2015eb91f321SVaclav Hapla #endif
201698921bdaSJacob 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);
20173ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2018eb91f321SVaclav Hapla }
2019eb91f321SVaclav Hapla 
2020d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatEqual_MPIDense(Mat A, Mat B, PetscBool *flag)
2021d71ae5a4SJacob Faibussowitsch {
20226e4ee0c6SHong Zhang   Mat_MPIDense *matB = (Mat_MPIDense *)B->data, *matA = (Mat_MPIDense *)A->data;
20236e4ee0c6SHong Zhang   Mat           a, b;
202490ace30eSBarry Smith 
20256e4ee0c6SHong Zhang   PetscFunctionBegin;
20266e4ee0c6SHong Zhang   a = matA->A;
20276e4ee0c6SHong Zhang   b = matB->A;
20282cf15c64SPierre Jolivet   PetscCall(MatEqual(a, b, flag));
2029462c564dSBarry Smith   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, flag, 1, MPIU_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
20303ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
20316e4ee0c6SHong Zhang }
203290ace30eSBarry Smith 
203366976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MatTransMatMult_MPIDense_MPIDense(void *data)
2034d71ae5a4SJacob Faibussowitsch {
20356718818eSStefano Zampini   Mat_TransMatMultDense *atb = (Mat_TransMatMultDense *)data;
2036baa3c1c6SHong Zhang 
2037baa3c1c6SHong Zhang   PetscFunctionBegin;
20389566063dSJacob Faibussowitsch   PetscCall(PetscFree2(atb->sendbuf, atb->recvcounts));
20399566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&atb->atb));
20409566063dSJacob Faibussowitsch   PetscCall(PetscFree(atb));
20413ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2042baa3c1c6SHong Zhang }
2043baa3c1c6SHong Zhang 
204466976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MatMatTransMult_MPIDense_MPIDense(void *data)
2045d71ae5a4SJacob Faibussowitsch {
20466718818eSStefano Zampini   Mat_MatTransMultDense *abt = (Mat_MatTransMultDense *)data;
2047cc48ffa7SToby Isaac 
2048cc48ffa7SToby Isaac   PetscFunctionBegin;
20499566063dSJacob Faibussowitsch   PetscCall(PetscFree2(abt->buf[0], abt->buf[1]));
20509566063dSJacob Faibussowitsch   PetscCall(PetscFree2(abt->recvcounts, abt->recvdispls));
20519566063dSJacob Faibussowitsch   PetscCall(PetscFree(abt));
20523ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2053cc48ffa7SToby Isaac }
2054cc48ffa7SToby Isaac 
2055d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2056d71ae5a4SJacob Faibussowitsch {
2057baa3c1c6SHong Zhang   Mat_MPIDense          *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
20586718818eSStefano Zampini   Mat_TransMatMultDense *atb;
2059cb20be35SHong Zhang   MPI_Comm               comm;
20606718818eSStefano Zampini   PetscMPIInt            size, *recvcounts;
20616718818eSStefano Zampini   PetscScalar           *carray, *sendbuf;
2062637a0070SStefano Zampini   const PetscScalar     *atbarray;
2063a1f9a5e6SStefano Zampini   PetscInt               i, cN = C->cmap->N, proc, k, j, lda;
2064e68c0b26SHong Zhang   const PetscInt        *ranges;
2065cb20be35SHong Zhang 
2066cb20be35SHong Zhang   PetscFunctionBegin;
20676718818eSStefano Zampini   MatCheckProduct(C, 3);
20685f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
20696718818eSStefano Zampini   atb        = (Mat_TransMatMultDense *)C->product->data;
20706718818eSStefano Zampini   recvcounts = atb->recvcounts;
20716718818eSStefano Zampini   sendbuf    = atb->sendbuf;
20726718818eSStefano Zampini 
20739566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
20749566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
2075e68c0b26SHong Zhang 
2076c5ef1628SHong Zhang   /* compute atbarray = aseq^T * bseq */
2077fb842aefSJose E. Roman   PetscCall(MatTransposeMatMult(a->A, b->A, atb->atb ? MAT_REUSE_MATRIX : MAT_INITIAL_MATRIX, PETSC_DETERMINE, &atb->atb));
2078cb20be35SHong Zhang 
20799566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRanges(C, &ranges));
2080cb20be35SHong Zhang 
2081660d5466SHong Zhang   /* arrange atbarray into sendbuf */
20829566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(atb->atb, &atbarray));
20839566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(atb->atb, &lda));
2084637a0070SStefano Zampini   for (proc = 0, k = 0; proc < size; proc++) {
2085baa3c1c6SHong Zhang     for (j = 0; j < cN; j++) {
2086a1f9a5e6SStefano Zampini       for (i = ranges[proc]; i < ranges[proc + 1]; i++) sendbuf[k++] = atbarray[i + j * lda];
2087cb20be35SHong Zhang     }
2088cb20be35SHong Zhang   }
20899566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(atb->atb, &atbarray));
2090637a0070SStefano Zampini 
2091c5ef1628SHong Zhang   /* sum all atbarray to local values of C */
20929566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(c->A, &carray));
20939566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Reduce_scatter(sendbuf, carray, recvcounts, MPIU_SCALAR, MPIU_SUM, comm));
20949566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(c->A, &carray));
209567af85e8SPierre Jolivet   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
20969566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
20979566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
20983ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2099cb20be35SHong Zhang }
2100cb20be35SHong Zhang 
2101d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2102d71ae5a4SJacob Faibussowitsch {
2103cb20be35SHong Zhang   MPI_Comm               comm;
2104baa3c1c6SHong Zhang   PetscMPIInt            size;
2105660d5466SHong Zhang   PetscInt               cm = A->cmap->n, cM, cN = B->cmap->N;
2106baa3c1c6SHong Zhang   Mat_TransMatMultDense *atb;
210747d993e7Ssuyashtn   PetscBool              cisdense = PETSC_FALSE;
21080acaeb71SStefano Zampini   const PetscInt        *ranges;
2109cb20be35SHong Zhang 
2110cb20be35SHong Zhang   PetscFunctionBegin;
21118f83a4d9SJacob Faibussowitsch   MatCheckProduct(C, 4);
21125f80ce2aSJacob Faibussowitsch   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
21139566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
2114cb20be35SHong Zhang   if (A->rmap->rstart != B->rmap->rstart || A->rmap->rend != B->rmap->rend) {
211598921bdaSJacob 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);
2116cb20be35SHong Zhang   }
2117cb20be35SHong Zhang 
21184222ddf1SHong Zhang   /* create matrix product C */
21199566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(C, cm, B->cmap->n, A->cmap->N, B->cmap->N));
212047d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA)
21219566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSECUDA, ""));
212247d993e7Ssuyashtn #endif
212347d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
212447d993e7Ssuyashtn   PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSEHIP, ""));
212547d993e7Ssuyashtn #endif
212648a46eb9SPierre Jolivet   if (!cisdense) PetscCall(MatSetType(C, ((PetscObject)A)->type_name));
21279566063dSJacob Faibussowitsch   PetscCall(MatSetUp(C));
2128baa3c1c6SHong Zhang 
21294222ddf1SHong Zhang   /* create data structure for reuse C */
21309566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
21319566063dSJacob Faibussowitsch   PetscCall(PetscNew(&atb));
21324222ddf1SHong Zhang   cM = C->rmap->N;
21339566063dSJacob Faibussowitsch   PetscCall(PetscMalloc2(cM * cN, &atb->sendbuf, size, &atb->recvcounts));
21349566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRanges(C, &ranges));
21356497c311SBarry Smith   for (PetscMPIInt i = 0; i < size; i++) PetscCall(PetscMPIIntCast((ranges[i + 1] - ranges[i]) * cN, &atb->recvcounts[i]));
21366718818eSStefano Zampini   C->product->data    = atb;
21376718818eSStefano Zampini   C->product->destroy = MatDestroy_MatTransMatMult_MPIDense_MPIDense;
21383ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2139cb20be35SHong Zhang }
2140cb20be35SHong Zhang 
2141d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2142d71ae5a4SJacob Faibussowitsch {
2143cc48ffa7SToby Isaac   MPI_Comm               comm;
2144cc48ffa7SToby Isaac   PetscMPIInt            i, size;
2145cc48ffa7SToby Isaac   PetscInt               maxRows, bufsiz;
2146cc48ffa7SToby Isaac   PetscMPIInt            tag;
21474222ddf1SHong Zhang   PetscInt               alg;
2148cc48ffa7SToby Isaac   Mat_MatTransMultDense *abt;
21494222ddf1SHong Zhang   Mat_Product           *product = C->product;
21504222ddf1SHong Zhang   PetscBool              flg;
2151cc48ffa7SToby Isaac 
2152cc48ffa7SToby Isaac   PetscFunctionBegin;
21536718818eSStefano Zampini   MatCheckProduct(C, 4);
21545f80ce2aSJacob Faibussowitsch   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
21554222ddf1SHong Zhang   /* check local size of A and B */
21565f80ce2aSJacob 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);
2157cc48ffa7SToby Isaac 
21589566063dSJacob Faibussowitsch   PetscCall(PetscStrcmp(product->alg, "allgatherv", &flg));
2159637a0070SStefano Zampini   alg = flg ? 0 : 1;
2160cc48ffa7SToby Isaac 
21614222ddf1SHong Zhang   /* setup matrix product C */
21629566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(C, A->rmap->n, B->rmap->n, A->rmap->N, B->rmap->N));
21639566063dSJacob Faibussowitsch   PetscCall(MatSetType(C, MATMPIDENSE));
21649566063dSJacob Faibussowitsch   PetscCall(MatSetUp(C));
21659566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag));
21664222ddf1SHong Zhang 
21674222ddf1SHong Zhang   /* create data structure for reuse C */
21689566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
21699566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
21709566063dSJacob Faibussowitsch   PetscCall(PetscNew(&abt));
2171cc48ffa7SToby Isaac   abt->tag = tag;
2172faa55883SToby Isaac   abt->alg = alg;
2173faa55883SToby Isaac   switch (alg) {
21744222ddf1SHong Zhang   case 1: /* alg: "cyclic" */
217557508eceSPierre Jolivet     for (maxRows = 0, i = 0; i < size; i++) maxRows = PetscMax(maxRows, B->rmap->range[i + 1] - B->rmap->range[i]);
2176cc48ffa7SToby Isaac     bufsiz = A->cmap->N * maxRows;
2177f4f49eeaSPierre Jolivet     PetscCall(PetscMalloc2(bufsiz, &abt->buf[0], bufsiz, &abt->buf[1]));
2178faa55883SToby Isaac     break;
21794222ddf1SHong Zhang   default: /* alg: "allgatherv" */
2180f4f49eeaSPierre Jolivet     PetscCall(PetscMalloc2(B->rmap->n * B->cmap->N, &abt->buf[0], B->rmap->N * B->cmap->N, &abt->buf[1]));
2181f4f49eeaSPierre Jolivet     PetscCall(PetscMalloc2(size, &abt->recvcounts, size + 1, &abt->recvdispls));
21826497c311SBarry Smith     for (i = 0; i <= size; i++) PetscCall(PetscMPIIntCast(B->rmap->range[i] * A->cmap->N, &abt->recvdispls[i]));
21836497c311SBarry Smith     for (i = 0; i < size; i++) PetscCall(PetscMPIIntCast(abt->recvdispls[i + 1] - abt->recvdispls[i], &abt->recvcounts[i]));
2184faa55883SToby Isaac     break;
2185faa55883SToby Isaac   }
2186cc48ffa7SToby Isaac 
21876718818eSStefano Zampini   C->product->data                = abt;
21886718818eSStefano Zampini   C->product->destroy             = MatDestroy_MatMatTransMult_MPIDense_MPIDense;
21894222ddf1SHong Zhang   C->ops->mattransposemultnumeric = MatMatTransposeMultNumeric_MPIDense_MPIDense;
21903ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2191cc48ffa7SToby Isaac }
2192cc48ffa7SToby Isaac 
2193d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(Mat A, Mat B, Mat C)
2194d71ae5a4SJacob Faibussowitsch {
2195cc48ffa7SToby Isaac   Mat_MPIDense          *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
21966718818eSStefano Zampini   Mat_MatTransMultDense *abt;
2197cc48ffa7SToby Isaac   MPI_Comm               comm;
21986497c311SBarry Smith   PetscMPIInt            rank, size, sendto, recvfrom, recvisfrom;
2199f4259b30SLisandro Dalcin   PetscScalar           *sendbuf, *recvbuf = NULL, *cv;
22006497c311SBarry Smith   PetscInt               i, cK             = A->cmap->N, sendsiz, recvsiz, k, j, bn;
2201cc48ffa7SToby Isaac   PetscScalar            _DOne = 1.0, _DZero = 0.0;
2202637a0070SStefano Zampini   const PetscScalar     *av, *bv;
220323fff9afSBarry Smith   PetscBLASInt           cm, cn, ck, alda, blda = 0, clda;
2204cc48ffa7SToby Isaac   MPI_Request            reqs[2];
2205cc48ffa7SToby Isaac   const PetscInt        *ranges;
2206cc48ffa7SToby Isaac 
2207cc48ffa7SToby Isaac   PetscFunctionBegin;
22086718818eSStefano Zampini   MatCheckProduct(C, 3);
22095f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
22106718818eSStefano Zampini   abt = (Mat_MatTransMultDense *)C->product->data;
22119566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
22129566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(comm, &rank));
22139566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
22149566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, &av));
22159566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(b->A, &bv));
22169566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(c->A, &cv));
22179566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &i));
22189566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &alda));
22199566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(b->A, &i));
22209566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &blda));
22219566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(c->A, &i));
22229566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &clda));
22239566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRanges(B, &ranges));
2224cc48ffa7SToby Isaac   bn = B->rmap->n;
2225637a0070SStefano Zampini   if (blda == bn) {
2226637a0070SStefano Zampini     sendbuf = (PetscScalar *)bv;
2227cc48ffa7SToby Isaac   } else {
2228cc48ffa7SToby Isaac     sendbuf = abt->buf[0];
2229cc48ffa7SToby Isaac     for (k = 0, i = 0; i < cK; i++) {
2230ad540459SPierre Jolivet       for (j = 0; j < bn; j++, k++) sendbuf[k] = bv[i * blda + j];
2231cc48ffa7SToby Isaac     }
2232cc48ffa7SToby Isaac   }
2233cc48ffa7SToby Isaac   if (size > 1) {
2234cc48ffa7SToby Isaac     sendto   = (rank + size - 1) % size;
2235cc48ffa7SToby Isaac     recvfrom = (rank + size + 1) % size;
2236cc48ffa7SToby Isaac   } else {
2237cc48ffa7SToby Isaac     sendto = recvfrom = 0;
2238cc48ffa7SToby Isaac   }
22399566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(cK, &ck));
22409566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm));
2241cc48ffa7SToby Isaac   recvisfrom = rank;
2242cc48ffa7SToby Isaac   for (i = 0; i < size; i++) {
2243cc48ffa7SToby Isaac     /* we have finished receiving in sending, bufs can be read/modified */
22446497c311SBarry Smith     PetscMPIInt nextrecvisfrom = (recvisfrom + 1) % size; /* which process the next recvbuf will originate on */
2245cc48ffa7SToby Isaac     PetscInt    nextbn         = ranges[nextrecvisfrom + 1] - ranges[nextrecvisfrom];
2246cc48ffa7SToby Isaac 
2247cc48ffa7SToby Isaac     if (nextrecvisfrom != rank) {
2248cc48ffa7SToby Isaac       /* start the cyclic sends from sendbuf, to recvbuf (which will switch to sendbuf) */
2249cc48ffa7SToby Isaac       sendsiz = cK * bn;
2250cc48ffa7SToby Isaac       recvsiz = cK * nextbn;
2251cc48ffa7SToby Isaac       recvbuf = (i & 1) ? abt->buf[0] : abt->buf[1];
22526497c311SBarry Smith       PetscCallMPI(MPIU_Isend(sendbuf, sendsiz, MPIU_SCALAR, sendto, abt->tag, comm, &reqs[0]));
22536497c311SBarry Smith       PetscCallMPI(MPIU_Irecv(recvbuf, recvsiz, MPIU_SCALAR, recvfrom, abt->tag, comm, &reqs[1]));
2254cc48ffa7SToby Isaac     }
2255cc48ffa7SToby Isaac 
2256cc48ffa7SToby Isaac     /* local aseq * sendbuf^T */
22579566063dSJacob Faibussowitsch     PetscCall(PetscBLASIntCast(ranges[recvisfrom + 1] - ranges[recvisfrom], &cn));
2258792fecdfSBarry Smith     if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "T", &cm, &cn, &ck, &_DOne, av, &alda, sendbuf, &cn, &_DZero, cv + clda * ranges[recvisfrom], &clda));
2259cc48ffa7SToby Isaac 
2260cc48ffa7SToby Isaac     if (nextrecvisfrom != rank) {
2261cc48ffa7SToby Isaac       /* wait for the sends and receives to complete, swap sendbuf and recvbuf */
22629566063dSJacob Faibussowitsch       PetscCallMPI(MPI_Waitall(2, reqs, MPI_STATUSES_IGNORE));
2263cc48ffa7SToby Isaac     }
2264cc48ffa7SToby Isaac     bn         = nextbn;
2265cc48ffa7SToby Isaac     recvisfrom = nextrecvisfrom;
2266cc48ffa7SToby Isaac     sendbuf    = recvbuf;
2267cc48ffa7SToby Isaac   }
22689566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, &av));
22699566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
22709566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(c->A, &cv));
227167af85e8SPierre Jolivet   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
22729566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
22739566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
22743ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2275cc48ffa7SToby Isaac }
2276cc48ffa7SToby Isaac 
2277d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(Mat A, Mat B, Mat C)
2278d71ae5a4SJacob Faibussowitsch {
2279faa55883SToby Isaac   Mat_MPIDense          *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
22806718818eSStefano Zampini   Mat_MatTransMultDense *abt;
2281faa55883SToby Isaac   MPI_Comm               comm;
22826497c311SBarry Smith   PetscMPIInt            size, ibn;
2283637a0070SStefano Zampini   PetscScalar           *cv, *sendbuf, *recvbuf;
2284637a0070SStefano Zampini   const PetscScalar     *av, *bv;
2285637a0070SStefano Zampini   PetscInt               blda, i, cK = A->cmap->N, k, j, bn;
2286faa55883SToby Isaac   PetscScalar            _DOne = 1.0, _DZero = 0.0;
2287637a0070SStefano Zampini   PetscBLASInt           cm, cn, ck, alda, clda;
2288faa55883SToby Isaac 
2289faa55883SToby Isaac   PetscFunctionBegin;
22906718818eSStefano Zampini   MatCheckProduct(C, 3);
22915f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
22926718818eSStefano Zampini   abt = (Mat_MatTransMultDense *)C->product->data;
22939566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
22949566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
22959566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, &av));
22969566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(b->A, &bv));
22979566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(c->A, &cv));
22989566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &i));
22999566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &alda));
23009566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(b->A, &blda));
23019566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(c->A, &i));
23029566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &clda));
2303faa55883SToby Isaac   /* copy transpose of B into buf[0] */
2304faa55883SToby Isaac   bn      = B->rmap->n;
2305faa55883SToby Isaac   sendbuf = abt->buf[0];
2306faa55883SToby Isaac   recvbuf = abt->buf[1];
2307faa55883SToby Isaac   for (k = 0, j = 0; j < bn; j++) {
2308ad540459SPierre Jolivet     for (i = 0; i < cK; i++, k++) sendbuf[k] = bv[i * blda + j];
2309faa55883SToby Isaac   }
23109566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
23116497c311SBarry Smith   PetscCall(PetscMPIIntCast(bn * cK, &ibn));
23126497c311SBarry Smith   PetscCallMPI(MPI_Allgatherv(sendbuf, ibn, MPIU_SCALAR, recvbuf, abt->recvcounts, abt->recvdispls, MPIU_SCALAR, comm));
23139566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(cK, &ck));
23149566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm));
23159566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(c->A->cmap->n, &cn));
2316792fecdfSBarry Smith   if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "N", &cm, &cn, &ck, &_DOne, av, &alda, recvbuf, &ck, &_DZero, cv, &clda));
23179566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, &av));
23189566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
23199566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(c->A, &cv));
232067af85e8SPierre Jolivet   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
23219566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
23229566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
23233ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2324faa55883SToby Isaac }
2325faa55883SToby Isaac 
2326d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2327d71ae5a4SJacob Faibussowitsch {
23286718818eSStefano Zampini   Mat_MatTransMultDense *abt;
2329faa55883SToby Isaac 
2330faa55883SToby Isaac   PetscFunctionBegin;
23316718818eSStefano Zampini   MatCheckProduct(C, 3);
23325f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
23336718818eSStefano Zampini   abt = (Mat_MatTransMultDense *)C->product->data;
2334faa55883SToby Isaac   switch (abt->alg) {
2335d71ae5a4SJacob Faibussowitsch   case 1:
2336d71ae5a4SJacob Faibussowitsch     PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(A, B, C));
2337d71ae5a4SJacob Faibussowitsch     break;
2338d71ae5a4SJacob Faibussowitsch   default:
2339d71ae5a4SJacob Faibussowitsch     PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(A, B, C));
2340d71ae5a4SJacob Faibussowitsch     break;
2341faa55883SToby Isaac   }
23423ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2343faa55883SToby Isaac }
2344faa55883SToby Isaac 
234566976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MatMatMult_MPIDense_MPIDense(void *data)
2346d71ae5a4SJacob Faibussowitsch {
23476718818eSStefano Zampini   Mat_MatMultDense *ab = (Mat_MatMultDense *)data;
2348320f2790SHong Zhang 
2349320f2790SHong Zhang   PetscFunctionBegin;
23509566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&ab->Ce));
23519566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&ab->Ae));
23529566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&ab->Be));
23539566063dSJacob Faibussowitsch   PetscCall(PetscFree(ab));
23543ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2355320f2790SHong Zhang }
2356320f2790SHong Zhang 
23575d8c7819SPierre Jolivet static PetscErrorCode MatMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2358d71ae5a4SJacob Faibussowitsch {
23596718818eSStefano Zampini   Mat_MatMultDense *ab;
23605d8c7819SPierre Jolivet   Mat_MPIDense     *mdn = (Mat_MPIDense *)A->data;
2361320f2790SHong Zhang 
2362320f2790SHong Zhang   PetscFunctionBegin;
23636718818eSStefano Zampini   MatCheckProduct(C, 3);
23645f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Missing product data");
23656718818eSStefano Zampini   ab = (Mat_MatMultDense *)C->product->data;
23665d8c7819SPierre Jolivet   if (ab->Ae && ab->Ce) {
23675d8c7819SPierre Jolivet #if PetscDefined(HAVE_ELEMENTAL)
23689566063dSJacob Faibussowitsch     PetscCall(MatConvert_MPIDense_Elemental(A, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Ae));
23699566063dSJacob Faibussowitsch     PetscCall(MatConvert_MPIDense_Elemental(B, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Be));
23709566063dSJacob Faibussowitsch     PetscCall(MatMatMultNumeric_Elemental(ab->Ae, ab->Be, ab->Ce));
23719566063dSJacob Faibussowitsch     PetscCall(MatConvert(ab->Ce, MATMPIDENSE, MAT_REUSE_MATRIX, &C));
23725d8c7819SPierre Jolivet #else
23735d8c7819SPierre Jolivet     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined");
23745d8c7819SPierre Jolivet #endif
23755d8c7819SPierre Jolivet   } else {
23765d8c7819SPierre Jolivet     const PetscScalar *read;
23775d8c7819SPierre Jolivet     PetscScalar       *write;
23785d8c7819SPierre Jolivet     PetscInt           lda;
23795d8c7819SPierre Jolivet 
23805d8c7819SPierre Jolivet     PetscCall(MatDenseGetLDA(B, &lda));
23815d8c7819SPierre Jolivet     PetscCall(MatDenseGetArrayRead(B, &read));
23825d8c7819SPierre Jolivet     PetscCall(MatDenseGetArrayWrite(ab->Be, &write));
23835d8c7819SPierre Jolivet     if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A)); /* cannot be done during the symbolic phase because of possible calls to MatProductReplaceMats() */
23845d8c7819SPierre Jolivet     for (PetscInt i = 0; i < C->cmap->N; ++i) {
23855d8c7819SPierre Jolivet       PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE));
23865d8c7819SPierre Jolivet       PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE));
23875d8c7819SPierre Jolivet     }
23885d8c7819SPierre Jolivet     PetscCall(MatDenseRestoreArrayWrite(ab->Be, &write));
23895d8c7819SPierre Jolivet     PetscCall(MatDenseRestoreArrayRead(B, &read));
2390f4f49eeaSPierre Jolivet     PetscCall(MatMatMultNumeric_SeqDense_SeqDense(((Mat_MPIDense *)A->data)->A, ab->Be, ((Mat_MPIDense *)C->data)->A));
23915d8c7819SPierre Jolivet   }
23923ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2393320f2790SHong Zhang }
2394320f2790SHong Zhang 
2395d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2396d71ae5a4SJacob Faibussowitsch {
23975d8c7819SPierre Jolivet   Mat_Product      *product = C->product;
23985d8c7819SPierre Jolivet   PetscInt          alg;
2399320f2790SHong Zhang   Mat_MatMultDense *ab;
24005d8c7819SPierre Jolivet   PetscBool         flg;
2401320f2790SHong Zhang 
2402320f2790SHong Zhang   PetscFunctionBegin;
24036718818eSStefano Zampini   MatCheckProduct(C, 4);
24045f80ce2aSJacob Faibussowitsch   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
24054222ddf1SHong Zhang   /* check local size of A and B */
24065d8c7819SPierre 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 ")",
24075d8c7819SPierre Jolivet              A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
2408320f2790SHong Zhang 
24095d8c7819SPierre Jolivet   PetscCall(PetscStrcmp(product->alg, "petsc", &flg));
24105d8c7819SPierre Jolivet   alg = flg ? 0 : 1;
24114222ddf1SHong Zhang 
24124222ddf1SHong Zhang   /* setup C */
24135d8c7819SPierre Jolivet   PetscCall(MatSetSizes(C, A->rmap->n, B->cmap->n, A->rmap->N, B->cmap->N));
24145d8c7819SPierre Jolivet   PetscCall(MatSetType(C, MATMPIDENSE));
24159566063dSJacob Faibussowitsch   PetscCall(MatSetUp(C));
2416320f2790SHong Zhang 
2417320f2790SHong Zhang   /* create data structure for reuse Cdense */
24189566063dSJacob Faibussowitsch   PetscCall(PetscNew(&ab));
24195d8c7819SPierre Jolivet 
24205d8c7819SPierre Jolivet   switch (alg) {
24215d8c7819SPierre Jolivet   case 1: /* alg: "elemental" */
24225d8c7819SPierre Jolivet #if PetscDefined(HAVE_ELEMENTAL)
24235d8c7819SPierre Jolivet     /* create elemental matrices Ae and Be */
24245d8c7819SPierre Jolivet     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &ab->Ae));
24255d8c7819SPierre Jolivet     PetscCall(MatSetSizes(ab->Ae, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N));
24265d8c7819SPierre Jolivet     PetscCall(MatSetType(ab->Ae, MATELEMENTAL));
24275d8c7819SPierre Jolivet     PetscCall(MatSetUp(ab->Ae));
24285d8c7819SPierre Jolivet     PetscCall(MatSetOption(ab->Ae, MAT_ROW_ORIENTED, PETSC_FALSE));
24295d8c7819SPierre Jolivet 
24305d8c7819SPierre Jolivet     PetscCall(MatCreate(PetscObjectComm((PetscObject)B), &ab->Be));
24315d8c7819SPierre Jolivet     PetscCall(MatSetSizes(ab->Be, PETSC_DECIDE, PETSC_DECIDE, B->rmap->N, B->cmap->N));
24325d8c7819SPierre Jolivet     PetscCall(MatSetType(ab->Be, MATELEMENTAL));
24335d8c7819SPierre Jolivet     PetscCall(MatSetUp(ab->Be));
24345d8c7819SPierre Jolivet     PetscCall(MatSetOption(ab->Be, MAT_ROW_ORIENTED, PETSC_FALSE));
24355d8c7819SPierre Jolivet 
24365d8c7819SPierre Jolivet     /* compute symbolic Ce = Ae*Be */
24375d8c7819SPierre Jolivet     PetscCall(MatCreate(PetscObjectComm((PetscObject)C), &ab->Ce));
24385d8c7819SPierre Jolivet     PetscCall(MatMatMultSymbolic_Elemental(ab->Ae, ab->Be, fill, ab->Ce));
24395d8c7819SPierre Jolivet #else
24405d8c7819SPierre Jolivet     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined");
24415d8c7819SPierre Jolivet #endif
24425d8c7819SPierre Jolivet     break;
24435d8c7819SPierre Jolivet   default: /* alg: "petsc" */
24445d8c7819SPierre Jolivet     ab->Ae = NULL;
24455d8c7819SPierre Jolivet     PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, A->cmap->N, B->cmap->N, NULL, &ab->Be));
24465d8c7819SPierre Jolivet     ab->Ce = NULL;
24475d8c7819SPierre Jolivet     break;
24485d8c7819SPierre Jolivet   }
24496718818eSStefano Zampini 
24506718818eSStefano Zampini   C->product->data       = ab;
24516718818eSStefano Zampini   C->product->destroy    = MatDestroy_MatMatMult_MPIDense_MPIDense;
24524222ddf1SHong Zhang   C->ops->matmultnumeric = MatMatMultNumeric_MPIDense_MPIDense;
24533ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2454320f2790SHong Zhang }
24552ef1f0ffSBarry Smith 
2456d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_AB(Mat C)
2457d71ae5a4SJacob Faibussowitsch {
24585d8c7819SPierre Jolivet   Mat_Product *product     = C->product;
24595d8c7819SPierre Jolivet   const char  *algTypes[2] = {"petsc", "elemental"};
24605d8c7819SPierre Jolivet   PetscInt     alg, nalg = PetscDefined(HAVE_ELEMENTAL) ? 2 : 1;
24615d8c7819SPierre Jolivet   PetscBool    flg = PETSC_FALSE;
24625d8c7819SPierre Jolivet 
24634d86920dSPierre Jolivet   PetscFunctionBegin;
24645d8c7819SPierre Jolivet   /* Set default algorithm */
24655d8c7819SPierre Jolivet   alg = 0; /* default is petsc */
24665d8c7819SPierre Jolivet   PetscCall(PetscStrcmp(product->alg, "default", &flg));
24675d8c7819SPierre Jolivet   if (flg) PetscCall(MatProductSetAlgorithm(C, (MatProductAlgorithm)algTypes[alg]));
24685d8c7819SPierre Jolivet 
24695d8c7819SPierre Jolivet   /* Get runtime option */
24705d8c7819SPierre Jolivet   PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_AB", "Mat");
24715d8c7819SPierre Jolivet   PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_AB", algTypes, nalg, algTypes[alg], &alg, &flg));
24725d8c7819SPierre Jolivet   PetscOptionsEnd();
24735d8c7819SPierre Jolivet   if (flg) PetscCall(MatProductSetAlgorithm(C, (MatProductAlgorithm)algTypes[alg]));
24745d8c7819SPierre Jolivet 
24754222ddf1SHong Zhang   C->ops->matmultsymbolic = MatMatMultSymbolic_MPIDense_MPIDense;
24764222ddf1SHong Zhang   C->ops->productsymbolic = MatProductSymbolic_AB;
24773ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2478320f2790SHong Zhang }
247986aefd0dSHong Zhang 
2480d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_AtB(Mat C)
2481d71ae5a4SJacob Faibussowitsch {
24824222ddf1SHong Zhang   Mat_Product *product = C->product;
24834222ddf1SHong Zhang   Mat          A = product->A, B = product->B;
24844222ddf1SHong Zhang 
24854222ddf1SHong Zhang   PetscFunctionBegin;
248607e90273SPierre 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 ")",
248707e90273SPierre Jolivet              A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
24886718818eSStefano Zampini   C->ops->transposematmultsymbolic = MatTransposeMatMultSymbolic_MPIDense_MPIDense;
24896718818eSStefano Zampini   C->ops->productsymbolic          = MatProductSymbolic_AtB;
24903ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
24914222ddf1SHong Zhang }
24924222ddf1SHong Zhang 
2493d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_ABt(Mat C)
2494d71ae5a4SJacob Faibussowitsch {
24954222ddf1SHong Zhang   Mat_Product *product     = C->product;
24964222ddf1SHong Zhang   const char  *algTypes[2] = {"allgatherv", "cyclic"};
24974222ddf1SHong Zhang   PetscInt     alg, nalg = 2;
24984222ddf1SHong Zhang   PetscBool    flg = PETSC_FALSE;
24994222ddf1SHong Zhang 
25004222ddf1SHong Zhang   PetscFunctionBegin;
25014222ddf1SHong Zhang   /* Set default algorithm */
25024222ddf1SHong Zhang   alg = 0; /* default is allgatherv */
25039566063dSJacob Faibussowitsch   PetscCall(PetscStrcmp(product->alg, "default", &flg));
250448a46eb9SPierre Jolivet   if (flg) PetscCall(MatProductSetAlgorithm(C, (MatProductAlgorithm)algTypes[alg]));
25054222ddf1SHong Zhang 
25064222ddf1SHong Zhang   /* Get runtime option */
25074222ddf1SHong Zhang   if (product->api_user) {
2508d0609cedSBarry Smith     PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatMatTransposeMult", "Mat");
25099566063dSJacob Faibussowitsch     PetscCall(PetscOptionsEList("-matmattransmult_mpidense_mpidense_via", "Algorithmic approach", "MatMatTransposeMult", algTypes, nalg, algTypes[alg], &alg, &flg));
2510d0609cedSBarry Smith     PetscOptionsEnd();
25114222ddf1SHong Zhang   } else {
2512d0609cedSBarry Smith     PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_ABt", "Mat");
25139566063dSJacob Faibussowitsch     PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_ABt", algTypes, nalg, algTypes[alg], &alg, &flg));
2514d0609cedSBarry Smith     PetscOptionsEnd();
25154222ddf1SHong Zhang   }
251648a46eb9SPierre Jolivet   if (flg) PetscCall(MatProductSetAlgorithm(C, (MatProductAlgorithm)algTypes[alg]));
25174222ddf1SHong Zhang 
25184222ddf1SHong Zhang   C->ops->mattransposemultsymbolic = MatMatTransposeMultSymbolic_MPIDense_MPIDense;
25194222ddf1SHong Zhang   C->ops->productsymbolic          = MatProductSymbolic_ABt;
25203ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
25214222ddf1SHong Zhang }
25224222ddf1SHong Zhang 
25235d8c7819SPierre Jolivet static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat C)
2524d71ae5a4SJacob Faibussowitsch {
25254222ddf1SHong Zhang   Mat_Product *product = C->product;
25264222ddf1SHong Zhang 
25274222ddf1SHong Zhang   PetscFunctionBegin;
25284222ddf1SHong Zhang   switch (product->type) {
2529d71ae5a4SJacob Faibussowitsch   case MATPRODUCT_AB:
2530d71ae5a4SJacob Faibussowitsch     PetscCall(MatProductSetFromOptions_MPIDense_AB(C));
2531d71ae5a4SJacob Faibussowitsch     break;
2532d71ae5a4SJacob Faibussowitsch   case MATPRODUCT_AtB:
2533d71ae5a4SJacob Faibussowitsch     PetscCall(MatProductSetFromOptions_MPIDense_AtB(C));
2534d71ae5a4SJacob Faibussowitsch     break;
2535d71ae5a4SJacob Faibussowitsch   case MATPRODUCT_ABt:
2536d71ae5a4SJacob Faibussowitsch     PetscCall(MatProductSetFromOptions_MPIDense_ABt(C));
2537d71ae5a4SJacob Faibussowitsch     break;
2538d71ae5a4SJacob Faibussowitsch   default:
2539d71ae5a4SJacob Faibussowitsch     break;
25404222ddf1SHong Zhang   }
25413ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
25424222ddf1SHong Zhang }
2543