xref: /petsc/src/mat/impls/dense/mpi/mpidense.c (revision d1a032db6cd7c39db5bfaa476c8e42d0c0ea531b)
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>
11d016bddeSToby Isaac #include <petscdevice.h>
12d016bddeSToby Isaac #include <petsc/private/deviceimpl.h>
138965ea79SLois Curfman McInnes 
14ab92ecdeSBarry Smith /*@
1511a5261eSBarry Smith   MatDenseGetLocalMatrix - For a `MATMPIDENSE` or `MATSEQDENSE` matrix returns the sequential
16ab92ecdeSBarry Smith   matrix that represents the operator. For sequential matrices it returns itself.
17ab92ecdeSBarry Smith 
18ab92ecdeSBarry Smith   Input Parameter:
192ef1f0ffSBarry Smith . A - the sequential or MPI `MATDENSE` matrix
20ab92ecdeSBarry Smith 
21ab92ecdeSBarry Smith   Output Parameter:
22ab92ecdeSBarry Smith . B - the inner matrix
23ab92ecdeSBarry Smith 
248e6c10adSSatish Balay   Level: intermediate
258e6c10adSSatish Balay 
261cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MATMPIDENSE`, `MATSEQDENSE`
27ab92ecdeSBarry Smith @*/
28d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetLocalMatrix(Mat A, Mat *B)
29d71ae5a4SJacob Faibussowitsch {
30ab92ecdeSBarry Smith   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
31ace3abfcSBarry Smith   PetscBool     flg;
32ab92ecdeSBarry Smith 
33ab92ecdeSBarry Smith   PetscFunctionBegin;
34d5ea218eSStefano Zampini   PetscValidHeaderSpecific(A, MAT_CLASSID, 1);
354f572ea9SToby Isaac   PetscAssertPointer(B, 2);
369566063dSJacob Faibussowitsch   PetscCall(PetscObjectBaseTypeCompare((PetscObject)A, MATMPIDENSE, &flg));
372205254eSKarl Rupp   if (flg) *B = mat->A;
38f140bc17SStefano Zampini   else {
399566063dSJacob Faibussowitsch     PetscCall(PetscObjectBaseTypeCompare((PetscObject)A, MATSEQDENSE, &flg));
4028b400f6SJacob Faibussowitsch     PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "Not for matrix type %s", ((PetscObject)A)->type_name);
41f140bc17SStefano Zampini     *B = A;
42f140bc17SStefano Zampini   }
433ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
44ab92ecdeSBarry Smith }
45ab92ecdeSBarry Smith 
4666976f2fSJacob Faibussowitsch static PetscErrorCode MatCopy_MPIDense(Mat A, Mat B, MatStructure s)
47d71ae5a4SJacob Faibussowitsch {
48a76f77c3SStefano Zampini   Mat_MPIDense *Amat = (Mat_MPIDense *)A->data;
49a76f77c3SStefano Zampini   Mat_MPIDense *Bmat = (Mat_MPIDense *)B->data;
50a76f77c3SStefano Zampini 
51a76f77c3SStefano Zampini   PetscFunctionBegin;
529566063dSJacob Faibussowitsch   PetscCall(MatCopy(Amat->A, Bmat->A, s));
533ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
54a76f77c3SStefano Zampini }
55a76f77c3SStefano Zampini 
56d71ae5a4SJacob Faibussowitsch PetscErrorCode MatShift_MPIDense(Mat A, PetscScalar alpha)
57d71ae5a4SJacob Faibussowitsch {
582f605a99SJose E. Roman   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
592f605a99SJose E. Roman   PetscInt      j, lda, rstart = A->rmap->rstart, rend = A->rmap->rend, rend2;
602f605a99SJose E. Roman   PetscScalar  *v;
612f605a99SJose E. Roman 
622f605a99SJose E. Roman   PetscFunctionBegin;
632f605a99SJose E. Roman   PetscCall(MatDenseGetArray(mat->A, &v));
642f605a99SJose E. Roman   PetscCall(MatDenseGetLDA(mat->A, &lda));
652f605a99SJose E. Roman   rend2 = PetscMin(rend, A->cmap->N);
662f605a99SJose E. Roman   if (rend2 > rstart) {
672f605a99SJose E. Roman     for (j = rstart; j < rend2; j++) v[j - rstart + j * lda] += alpha;
682f605a99SJose E. Roman     PetscCall(PetscLogFlops(rend2 - rstart));
692f605a99SJose E. Roman   }
702f605a99SJose E. Roman   PetscCall(MatDenseRestoreArray(mat->A, &v));
713ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
722f605a99SJose E. Roman }
732f605a99SJose E. Roman 
7466976f2fSJacob Faibussowitsch static PetscErrorCode MatGetRow_MPIDense(Mat A, PetscInt row, PetscInt *nz, PetscInt **idx, PetscScalar **v)
75d71ae5a4SJacob Faibussowitsch {
76ba8c8a56SBarry Smith   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
77d0f46423SBarry Smith   PetscInt      lrow, rstart = A->rmap->rstart, rend = A->rmap->rend;
78ba8c8a56SBarry Smith 
79ba8c8a56SBarry Smith   PetscFunctionBegin;
80aed4548fSBarry Smith   PetscCheck(row >= rstart && row < rend, PETSC_COMM_SELF, PETSC_ERR_SUP, "only local rows");
81ba8c8a56SBarry Smith   lrow = row - rstart;
829566063dSJacob Faibussowitsch   PetscCall(MatGetRow(mat->A, lrow, nz, (const PetscInt **)idx, (const PetscScalar **)v));
833ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
84ba8c8a56SBarry Smith }
85ba8c8a56SBarry Smith 
8666976f2fSJacob Faibussowitsch static PetscErrorCode MatRestoreRow_MPIDense(Mat A, PetscInt row, PetscInt *nz, PetscInt **idx, PetscScalar **v)
87d71ae5a4SJacob Faibussowitsch {
88637a0070SStefano Zampini   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
89637a0070SStefano Zampini   PetscInt      lrow, rstart = A->rmap->rstart, rend = A->rmap->rend;
90ba8c8a56SBarry Smith 
91ba8c8a56SBarry Smith   PetscFunctionBegin;
92aed4548fSBarry Smith   PetscCheck(row >= rstart && row < rend, PETSC_COMM_SELF, PETSC_ERR_SUP, "only local rows");
93637a0070SStefano Zampini   lrow = row - rstart;
949566063dSJacob Faibussowitsch   PetscCall(MatRestoreRow(mat->A, lrow, nz, (const PetscInt **)idx, (const PetscScalar **)v));
953ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
96ba8c8a56SBarry Smith }
97ba8c8a56SBarry Smith 
9866976f2fSJacob Faibussowitsch static PetscErrorCode MatGetDiagonalBlock_MPIDense(Mat A, Mat *a)
99d71ae5a4SJacob Faibussowitsch {
1000de54da6SSatish Balay   Mat_MPIDense *mdn = (Mat_MPIDense *)A->data;
101d0f46423SBarry Smith   PetscInt      m = A->rmap->n, rstart = A->rmap->rstart;
10287828ca2SBarry Smith   PetscScalar  *array;
1030de54da6SSatish Balay   MPI_Comm      comm;
1044b6ae80fSStefano Zampini   PetscBool     flg;
10511bd1e4dSLisandro Dalcin   Mat           B;
1060de54da6SSatish Balay 
1070de54da6SSatish Balay   PetscFunctionBegin;
1089566063dSJacob Faibussowitsch   PetscCall(MatHasCongruentLayouts(A, &flg));
10928b400f6SJacob Faibussowitsch   PetscCheck(flg, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only square matrices supported.");
1109566063dSJacob Faibussowitsch   PetscCall(PetscObjectQuery((PetscObject)A, "DiagonalBlock", (PetscObject *)&B));
111616b8fbbSStefano Zampini   if (!B) { /* This should use MatDenseGetSubMatrix (not create), but we would need a call like MatRestoreDiagonalBlock */
112a51cd166SPierre Jolivet #if PetscDefined(HAVE_CUDA)
1139566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)mdn->A, MATSEQDENSECUDA, &flg));
11428b400f6SJacob 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);
115a51cd166SPierre Jolivet #elif PetscDefined(HAVE_HIP)
11647d993e7Ssuyashtn     PetscCall(PetscObjectTypeCompare((PetscObject)mdn->A, MATSEQDENSEHIP, &flg));
11747d993e7Ssuyashtn     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);
11847d993e7Ssuyashtn #endif
119f4f49eeaSPierre Jolivet     PetscCall(PetscObjectGetComm((PetscObject)mdn->A, &comm));
1209566063dSJacob Faibussowitsch     PetscCall(MatCreate(comm, &B));
1219566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(B, m, m, m, m));
1229566063dSJacob Faibussowitsch     PetscCall(MatSetType(B, ((PetscObject)mdn->A)->type_name));
1239566063dSJacob Faibussowitsch     PetscCall(MatDenseGetArrayRead(mdn->A, (const PetscScalar **)&array));
1249566063dSJacob Faibussowitsch     PetscCall(MatSeqDenseSetPreallocation(B, array + m * rstart));
1259566063dSJacob Faibussowitsch     PetscCall(MatDenseRestoreArrayRead(mdn->A, (const PetscScalar **)&array));
1269566063dSJacob Faibussowitsch     PetscCall(PetscObjectCompose((PetscObject)A, "DiagonalBlock", (PetscObject)B));
12711bd1e4dSLisandro Dalcin     *a = B;
1289566063dSJacob Faibussowitsch     PetscCall(MatDestroy(&B));
1292205254eSKarl Rupp   } else *a = B;
1303ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1310de54da6SSatish Balay }
1320de54da6SSatish Balay 
13366976f2fSJacob Faibussowitsch static PetscErrorCode MatSetValues_MPIDense(Mat mat, PetscInt m, const PetscInt idxm[], PetscInt n, const PetscInt idxn[], const PetscScalar v[], InsertMode addv)
134d71ae5a4SJacob Faibussowitsch {
13539b7565bSBarry Smith   Mat_MPIDense *A = (Mat_MPIDense *)mat->data;
136d0f46423SBarry Smith   PetscInt      i, j, rstart = mat->rmap->rstart, rend = mat->rmap->rend, row;
137ace3abfcSBarry Smith   PetscBool     roworiented = A->roworiented;
1388965ea79SLois Curfman McInnes 
1393a40ed3dSBarry Smith   PetscFunctionBegin;
1408965ea79SLois Curfman McInnes   for (i = 0; i < m; i++) {
1415ef9f2a5SBarry Smith     if (idxm[i] < 0) continue;
14208401ef6SPierre Jolivet     PetscCheck(idxm[i] < mat->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large");
1438965ea79SLois Curfman McInnes     if (idxm[i] >= rstart && idxm[i] < rend) {
1448965ea79SLois Curfman McInnes       row = idxm[i] - rstart;
14539b7565bSBarry Smith       if (roworiented) {
146de0a9f35SPierre Jolivet         PetscCall(MatSetValues(A->A, 1, &row, n, idxn, PetscSafePointerPlusOffset(v, i * n), addv));
1473a40ed3dSBarry Smith       } else {
1488965ea79SLois Curfman McInnes         for (j = 0; j < n; j++) {
1495ef9f2a5SBarry Smith           if (idxn[j] < 0) continue;
15008401ef6SPierre Jolivet           PetscCheck(idxn[j] < mat->cmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Column too large");
151de0a9f35SPierre Jolivet           PetscCall(MatSetValues(A->A, 1, &row, 1, &idxn[j], PetscSafePointerPlusOffset(v, i + j * m), addv));
15239b7565bSBarry Smith         }
1538965ea79SLois Curfman McInnes       }
1542205254eSKarl Rupp     } else if (!A->donotstash) {
1555080c13bSMatthew G Knepley       mat->assembled = PETSC_FALSE;
15639b7565bSBarry Smith       if (roworiented) {
157de0a9f35SPierre Jolivet         PetscCall(MatStashValuesRow_Private(&mat->stash, idxm[i], n, idxn, PetscSafePointerPlusOffset(v, i * n), PETSC_FALSE));
158d36fbae8SSatish Balay       } else {
159de0a9f35SPierre Jolivet         PetscCall(MatStashValuesCol_Private(&mat->stash, idxm[i], n, idxn, PetscSafePointerPlusOffset(v, i), m, PETSC_FALSE));
16039b7565bSBarry Smith       }
161b49de8d1SLois Curfman McInnes     }
162b49de8d1SLois Curfman McInnes   }
1633ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
164b49de8d1SLois Curfman McInnes }
165b49de8d1SLois Curfman McInnes 
16666976f2fSJacob Faibussowitsch static PetscErrorCode MatGetValues_MPIDense(Mat mat, PetscInt m, const PetscInt idxm[], PetscInt n, const PetscInt idxn[], PetscScalar v[])
167d71ae5a4SJacob Faibussowitsch {
168b49de8d1SLois Curfman McInnes   Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
169d0f46423SBarry Smith   PetscInt      i, j, rstart = mat->rmap->rstart, rend = mat->rmap->rend, row;
170b49de8d1SLois Curfman McInnes 
1713a40ed3dSBarry Smith   PetscFunctionBegin;
172b49de8d1SLois Curfman McInnes   for (i = 0; i < m; i++) {
17354c59aa7SJacob Faibussowitsch     if (idxm[i] < 0) continue; /* negative row */
17454c59aa7SJacob Faibussowitsch     PetscCheck(idxm[i] < mat->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large");
175966bd95aSPierre Jolivet     PetscCheck(idxm[i] >= rstart && idxm[i] < rend, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only local values currently supported");
176b49de8d1SLois Curfman McInnes     row = idxm[i] - rstart;
177b49de8d1SLois Curfman McInnes     for (j = 0; j < n; j++) {
17854c59aa7SJacob Faibussowitsch       if (idxn[j] < 0) continue; /* negative column */
17954c59aa7SJacob Faibussowitsch       PetscCheck(idxn[j] < mat->cmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Column too large");
1809566063dSJacob Faibussowitsch       PetscCall(MatGetValues(mdn->A, 1, &row, 1, &idxn[j], v + i * n + j));
181b49de8d1SLois Curfman McInnes     }
1828965ea79SLois Curfman McInnes   }
1833ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1848965ea79SLois Curfman McInnes }
1858965ea79SLois Curfman McInnes 
186d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetLDA_MPIDense(Mat A, PetscInt *lda)
187d71ae5a4SJacob Faibussowitsch {
18849a6ff4bSBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
18949a6ff4bSBarry Smith 
19049a6ff4bSBarry Smith   PetscFunctionBegin;
1919566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, lda));
1923ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
19349a6ff4bSBarry Smith }
19449a6ff4bSBarry Smith 
195d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseSetLDA_MPIDense(Mat A, PetscInt lda)
196d71ae5a4SJacob Faibussowitsch {
197ad16ce7aSStefano Zampini   Mat_MPIDense *a     = (Mat_MPIDense *)A->data;
19847d993e7Ssuyashtn   MatType       mtype = MATSEQDENSE;
199ad16ce7aSStefano Zampini 
200ad16ce7aSStefano Zampini   PetscFunctionBegin;
20117359960SJose E. Roman   if (!a->A) {
20228b400f6SJacob Faibussowitsch     PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2039566063dSJacob Faibussowitsch     PetscCall(PetscLayoutSetUp(A->rmap));
2049566063dSJacob Faibussowitsch     PetscCall(PetscLayoutSetUp(A->cmap));
2059566063dSJacob Faibussowitsch     PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
2069566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(a->A, A->rmap->n, A->cmap->N, A->rmap->n, A->cmap->N));
207a51cd166SPierre Jolivet #if PetscDefined(HAVE_CUDA)
20847d993e7Ssuyashtn     PetscBool iscuda;
2099566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSECUDA, &iscuda));
21047d993e7Ssuyashtn     if (iscuda) mtype = MATSEQDENSECUDA;
211a51cd166SPierre Jolivet #elif PetscDefined(HAVE_HIP)
21247d993e7Ssuyashtn     PetscBool iship;
21347d993e7Ssuyashtn     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSEHIP, &iship));
21447d993e7Ssuyashtn     if (iship) mtype = MATSEQDENSEHIP;
21547d993e7Ssuyashtn #endif
21647d993e7Ssuyashtn     PetscCall(MatSetType(a->A, mtype));
21717359960SJose E. Roman   }
2189566063dSJacob Faibussowitsch   PetscCall(MatDenseSetLDA(a->A, lda));
2193ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
220ad16ce7aSStefano Zampini }
221ad16ce7aSStefano Zampini 
222d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetArray_MPIDense(Mat A, PetscScalar **array)
223d71ae5a4SJacob Faibussowitsch {
224ff14e315SSatish Balay   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
225ff14e315SSatish Balay 
2263a40ed3dSBarry Smith   PetscFunctionBegin;
22728b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2289566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(a->A, array));
2293ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
230ff14e315SSatish Balay }
231ff14e315SSatish Balay 
232f1e99121SPierre Jolivet static PetscErrorCode MatDenseGetArrayRead_MPIDense(Mat A, PetscScalar **array)
233d71ae5a4SJacob Faibussowitsch {
2348572280aSBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
2358572280aSBarry Smith 
2368572280aSBarry Smith   PetscFunctionBegin;
23728b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
238f1e99121SPierre Jolivet   PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)array));
2393ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2408572280aSBarry Smith }
2418572280aSBarry Smith 
242d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetArrayWrite_MPIDense(Mat A, PetscScalar **array)
243d71ae5a4SJacob Faibussowitsch {
2446947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
2456947451fSStefano Zampini 
2466947451fSStefano Zampini   PetscFunctionBegin;
24728b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2489566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(a->A, array));
2493ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2506947451fSStefano Zampini }
2516947451fSStefano Zampini 
252d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDensePlaceArray_MPIDense(Mat A, const PetscScalar *array)
253d71ae5a4SJacob Faibussowitsch {
254d3042a70SBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
255d3042a70SBarry Smith 
256d3042a70SBarry Smith   PetscFunctionBegin;
25728b400f6SJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
25828b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2599566063dSJacob Faibussowitsch   PetscCall(MatDensePlaceArray(a->A, array));
2603ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
261d3042a70SBarry Smith }
262d3042a70SBarry Smith 
263d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseResetArray_MPIDense(Mat A)
264d71ae5a4SJacob Faibussowitsch {
265d3042a70SBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
266d3042a70SBarry Smith 
267d3042a70SBarry Smith   PetscFunctionBegin;
26828b400f6SJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
26928b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2709566063dSJacob Faibussowitsch   PetscCall(MatDenseResetArray(a->A));
2713ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
272d3042a70SBarry Smith }
273d3042a70SBarry Smith 
274d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseReplaceArray_MPIDense(Mat A, const PetscScalar *array)
275d71ae5a4SJacob Faibussowitsch {
276d5ea218eSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
277d5ea218eSStefano Zampini 
278d5ea218eSStefano Zampini   PetscFunctionBegin;
27928b400f6SJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
28028b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2819566063dSJacob Faibussowitsch   PetscCall(MatDenseReplaceArray(a->A, array));
2823ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
283d5ea218eSStefano Zampini }
284d5ea218eSStefano Zampini 
285d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatCreateSubMatrix_MPIDense(Mat A, IS isrow, IS iscol, MatReuse scall, Mat *B)
286d71ae5a4SJacob Faibussowitsch {
287ca3fa75bSLois Curfman McInnes   Mat_MPIDense      *mat = (Mat_MPIDense *)A->data, *newmatd;
288637a0070SStefano Zampini   PetscInt           lda, i, j, rstart, rend, nrows, ncols, Ncols, nlrows, nlcols;
2895d0c19d7SBarry Smith   const PetscInt    *irow, *icol;
290637a0070SStefano Zampini   const PetscScalar *v;
291637a0070SStefano Zampini   PetscScalar       *bv;
292ca3fa75bSLois Curfman McInnes   Mat                newmat;
2934aa3045dSJed Brown   IS                 iscol_local;
29442a884f0SBarry Smith   MPI_Comm           comm_is, comm_mat;
295ca3fa75bSLois Curfman McInnes 
296ca3fa75bSLois Curfman McInnes   PetscFunctionBegin;
2979566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm_mat));
2989566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)iscol, &comm_is));
29908401ef6SPierre Jolivet   PetscCheck(comm_mat == comm_is, PETSC_COMM_SELF, PETSC_ERR_ARG_NOTSAMECOMM, "IS communicator must match matrix communicator");
30042a884f0SBarry Smith 
3019566063dSJacob Faibussowitsch   PetscCall(ISAllGather(iscol, &iscol_local));
3029566063dSJacob Faibussowitsch   PetscCall(ISGetIndices(isrow, &irow));
3039566063dSJacob Faibussowitsch   PetscCall(ISGetIndices(iscol_local, &icol));
3049566063dSJacob Faibussowitsch   PetscCall(ISGetLocalSize(isrow, &nrows));
3059566063dSJacob Faibussowitsch   PetscCall(ISGetLocalSize(iscol, &ncols));
3069566063dSJacob Faibussowitsch   PetscCall(ISGetSize(iscol, &Ncols)); /* global number of columns, size of iscol_local */
307ca3fa75bSLois Curfman McInnes 
308ca3fa75bSLois Curfman McInnes   /* No parallel redistribution currently supported! Should really check each index set
309da81f932SPierre Jolivet      to confirm that it is OK.  ... Currently supports only submatrix same partitioning as
3107eba5e9cSLois Curfman McInnes      original matrix! */
311ca3fa75bSLois Curfman McInnes 
3129566063dSJacob Faibussowitsch   PetscCall(MatGetLocalSize(A, &nlrows, &nlcols));
3139566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(A, &rstart, &rend));
314ca3fa75bSLois Curfman McInnes 
315ca3fa75bSLois Curfman McInnes   /* Check submatrix call */
316ca3fa75bSLois Curfman McInnes   if (scall == MAT_REUSE_MATRIX) {
317e32f2f54SBarry Smith     /* SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Reused submatrix wrong size"); */
3187eba5e9cSLois Curfman McInnes     /* Really need to test rows and column sizes! */
319ca3fa75bSLois Curfman McInnes     newmat = *B;
320ca3fa75bSLois Curfman McInnes   } else {
321ca3fa75bSLois Curfman McInnes     /* Create and fill new matrix */
3229566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &newmat));
3239566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(newmat, nrows, ncols, PETSC_DECIDE, Ncols));
3249566063dSJacob Faibussowitsch     PetscCall(MatSetType(newmat, ((PetscObject)A)->type_name));
3259566063dSJacob Faibussowitsch     PetscCall(MatMPIDenseSetPreallocation(newmat, NULL));
326ca3fa75bSLois Curfman McInnes   }
327ca3fa75bSLois Curfman McInnes 
328ca3fa75bSLois Curfman McInnes   /* Now extract the data pointers and do the copy, column at a time */
329ca3fa75bSLois Curfman McInnes   newmatd = (Mat_MPIDense *)newmat->data;
3309566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(newmatd->A, &bv));
3319566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(mat->A, &v));
3329566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(mat->A, &lda));
3334aa3045dSJed Brown   for (i = 0; i < Ncols; i++) {
334637a0070SStefano Zampini     const PetscScalar *av = v + lda * icol[i];
335ad540459SPierre Jolivet     for (j = 0; j < nrows; j++) *bv++ = av[irow[j] - rstart];
336ca3fa75bSLois Curfman McInnes   }
3379566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(mat->A, &v));
3389566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(newmatd->A, &bv));
339ca3fa75bSLois Curfman McInnes 
340ca3fa75bSLois Curfman McInnes   /* Assemble the matrices so that the correct flags are set */
3419566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(newmat, MAT_FINAL_ASSEMBLY));
3429566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(newmat, MAT_FINAL_ASSEMBLY));
343ca3fa75bSLois Curfman McInnes 
344ca3fa75bSLois Curfman McInnes   /* Free work space */
3459566063dSJacob Faibussowitsch   PetscCall(ISRestoreIndices(isrow, &irow));
3469566063dSJacob Faibussowitsch   PetscCall(ISRestoreIndices(iscol_local, &icol));
3479566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&iscol_local));
348ca3fa75bSLois Curfman McInnes   *B = newmat;
3493ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
350ca3fa75bSLois Curfman McInnes }
351ca3fa75bSLois Curfman McInnes 
35266976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreArray_MPIDense(Mat A, PetscScalar **array)
353d71ae5a4SJacob Faibussowitsch {
35473a71a0fSBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
35573a71a0fSBarry Smith 
3563a40ed3dSBarry Smith   PetscFunctionBegin;
3579566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(a->A, array));
3583ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
359ff14e315SSatish Balay }
360ff14e315SSatish Balay 
361f1e99121SPierre Jolivet static PetscErrorCode MatDenseRestoreArrayRead_MPIDense(Mat A, PetscScalar **array)
362d71ae5a4SJacob Faibussowitsch {
3638572280aSBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
3648572280aSBarry Smith 
3658572280aSBarry Smith   PetscFunctionBegin;
366f1e99121SPierre Jolivet   PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)array));
3673ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3688572280aSBarry Smith }
3698572280aSBarry Smith 
37066976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreArrayWrite_MPIDense(Mat A, PetscScalar **array)
371d71ae5a4SJacob Faibussowitsch {
3726947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
3736947451fSStefano Zampini 
3746947451fSStefano Zampini   PetscFunctionBegin;
3759566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(a->A, array));
3763ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3776947451fSStefano Zampini }
3786947451fSStefano Zampini 
37966976f2fSJacob Faibussowitsch static PetscErrorCode MatAssemblyBegin_MPIDense(Mat mat, MatAssemblyType mode)
380d71ae5a4SJacob Faibussowitsch {
38139ddd567SLois Curfman McInnes   Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
38213f74950SBarry Smith   PetscInt      nstash, reallocs;
3838965ea79SLois Curfman McInnes 
3843a40ed3dSBarry Smith   PetscFunctionBegin;
3853ba16761SJacob Faibussowitsch   if (mdn->donotstash || mat->nooffprocentries) PetscFunctionReturn(PETSC_SUCCESS);
3868965ea79SLois Curfman McInnes 
3879566063dSJacob Faibussowitsch   PetscCall(MatStashScatterBegin_Private(mat, &mat->stash, mat->rmap->range));
3889566063dSJacob Faibussowitsch   PetscCall(MatStashGetInfo_Private(&mat->stash, &nstash, &reallocs));
3899566063dSJacob Faibussowitsch   PetscCall(PetscInfo(mdn->A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
3903ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3918965ea79SLois Curfman McInnes }
3928965ea79SLois Curfman McInnes 
39366976f2fSJacob Faibussowitsch static PetscErrorCode MatAssemblyEnd_MPIDense(Mat mat, MatAssemblyType mode)
394d71ae5a4SJacob Faibussowitsch {
39539ddd567SLois Curfman McInnes   Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
39613f74950SBarry Smith   PetscInt      i, *row, *col, flg, j, rstart, ncols;
39713f74950SBarry Smith   PetscMPIInt   n;
39887828ca2SBarry Smith   PetscScalar  *val;
3998965ea79SLois Curfman McInnes 
4003a40ed3dSBarry Smith   PetscFunctionBegin;
401910cf402Sprj-   if (!mdn->donotstash && !mat->nooffprocentries) {
4028965ea79SLois Curfman McInnes     /*  wait on receives */
4037ef1d9bdSSatish Balay     while (1) {
4049566063dSJacob Faibussowitsch       PetscCall(MatStashScatterGetMesg_Private(&mat->stash, &n, &row, &col, &val, &flg));
4057ef1d9bdSSatish Balay       if (!flg) break;
4068965ea79SLois Curfman McInnes 
4077ef1d9bdSSatish Balay       for (i = 0; i < n;) {
4087ef1d9bdSSatish Balay         /* Now identify the consecutive vals belonging to the same row */
4092205254eSKarl Rupp         for (j = i, rstart = row[j]; j < n; j++) {
4102205254eSKarl Rupp           if (row[j] != rstart) break;
4112205254eSKarl Rupp         }
4127ef1d9bdSSatish Balay         if (j < n) ncols = j - i;
4137ef1d9bdSSatish Balay         else ncols = n - i;
4147ef1d9bdSSatish Balay         /* Now assemble all these values with a single function call */
4159566063dSJacob Faibussowitsch         PetscCall(MatSetValues_MPIDense(mat, 1, row + i, ncols, col + i, val + i, mat->insertmode));
4167ef1d9bdSSatish Balay         i = j;
4178965ea79SLois Curfman McInnes       }
4187ef1d9bdSSatish Balay     }
4199566063dSJacob Faibussowitsch     PetscCall(MatStashScatterEnd_Private(&mat->stash));
420910cf402Sprj-   }
4218965ea79SLois Curfman McInnes 
4229566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(mdn->A, mode));
4239566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(mdn->A, mode));
4243ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4258965ea79SLois Curfman McInnes }
4268965ea79SLois Curfman McInnes 
42766976f2fSJacob Faibussowitsch static PetscErrorCode MatZeroEntries_MPIDense(Mat A)
428d71ae5a4SJacob Faibussowitsch {
42939ddd567SLois Curfman McInnes   Mat_MPIDense *l = (Mat_MPIDense *)A->data;
4303a40ed3dSBarry Smith 
4313a40ed3dSBarry Smith   PetscFunctionBegin;
4329566063dSJacob Faibussowitsch   PetscCall(MatZeroEntries(l->A));
4333ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4348965ea79SLois Curfman McInnes }
4358965ea79SLois Curfman McInnes 
43666976f2fSJacob Faibussowitsch static PetscErrorCode MatZeroRows_MPIDense(Mat A, PetscInt n, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
437d71ae5a4SJacob Faibussowitsch {
43839ddd567SLois Curfman McInnes   Mat_MPIDense *l = (Mat_MPIDense *)A->data;
439637a0070SStefano Zampini   PetscInt      i, len, *lrows;
440637a0070SStefano Zampini 
441637a0070SStefano Zampini   PetscFunctionBegin;
442637a0070SStefano Zampini   /* get locally owned rows */
4439566063dSJacob Faibussowitsch   PetscCall(PetscLayoutMapLocal(A->rmap, n, rows, &len, &lrows, NULL));
444dd8e379bSPierre Jolivet   /* fix right-hand side if needed */
445637a0070SStefano Zampini   if (x && b) {
44697b48c8fSBarry Smith     const PetscScalar *xx;
44797b48c8fSBarry Smith     PetscScalar       *bb;
4488965ea79SLois Curfman McInnes 
4499566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(x, &xx));
4509566063dSJacob Faibussowitsch     PetscCall(VecGetArrayWrite(b, &bb));
451637a0070SStefano Zampini     for (i = 0; i < len; ++i) bb[lrows[i]] = diag * xx[lrows[i]];
4529566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(x, &xx));
4539566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayWrite(b, &bb));
45497b48c8fSBarry Smith   }
4559566063dSJacob Faibussowitsch   PetscCall(MatZeroRows(l->A, len, lrows, 0.0, NULL, NULL));
456e2eb51b1SBarry Smith   if (diag != 0.0) {
457637a0070SStefano Zampini     Vec d;
458b9679d65SBarry Smith 
4599566063dSJacob Faibussowitsch     PetscCall(MatCreateVecs(A, NULL, &d));
4609566063dSJacob Faibussowitsch     PetscCall(VecSet(d, diag));
4619566063dSJacob Faibussowitsch     PetscCall(MatDiagonalSet(A, d, INSERT_VALUES));
4629566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&d));
463b9679d65SBarry Smith   }
4649566063dSJacob Faibussowitsch   PetscCall(PetscFree(lrows));
4653ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4668965ea79SLois Curfman McInnes }
4678965ea79SLois Curfman McInnes 
4680be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMult_SeqDense(Mat, Vec, Vec);
4690be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMultAdd_SeqDense(Mat, Vec, Vec, Vec);
4700be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMultTranspose_SeqDense(Mat, Vec, Vec);
4710be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMultTransposeAdd_SeqDense(Mat, Vec, Vec, Vec);
4720be0d8bdSHansol Suh 
473d016bddeSToby Isaac static PetscErrorCode MatMultColumnRange_MPIDense(Mat mat, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end)
474d016bddeSToby Isaac {
475d016bddeSToby Isaac   Mat_MPIDense      *mdn = (Mat_MPIDense *)mat->data;
476d016bddeSToby Isaac   const PetscScalar *ax;
477d016bddeSToby Isaac   PetscScalar       *ay;
478d016bddeSToby Isaac   PetscMemType       axmtype, aymtype;
479d016bddeSToby Isaac 
480d016bddeSToby Isaac   PetscFunctionBegin;
481d016bddeSToby Isaac   if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
482d016bddeSToby Isaac   PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
483d016bddeSToby Isaac   PetscCall(VecGetArrayWriteAndMemType(mdn->lvec, &ay, &aymtype));
484d016bddeSToby Isaac   PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
485d016bddeSToby Isaac   PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
486d016bddeSToby Isaac   PetscCall(VecRestoreArrayWriteAndMemType(mdn->lvec, &ay));
487d016bddeSToby Isaac   PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
488d016bddeSToby Isaac   PetscUseMethod(mdn->A, "MatMultColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (mdn->A, mdn->lvec, yy, c_start, c_end));
489d016bddeSToby Isaac   PetscFunctionReturn(PETSC_SUCCESS);
490d016bddeSToby Isaac }
491d016bddeSToby Isaac 
49266976f2fSJacob Faibussowitsch static PetscErrorCode MatMult_MPIDense(Mat mat, Vec xx, Vec yy)
493d71ae5a4SJacob Faibussowitsch {
49439ddd567SLois Curfman McInnes   Mat_MPIDense      *mdn = (Mat_MPIDense *)mat->data;
495637a0070SStefano Zampini   const PetscScalar *ax;
496637a0070SStefano Zampini   PetscScalar       *ay;
497d0295fc0SJunchao Zhang   PetscMemType       axmtype, aymtype;
498c456f294SBarry Smith 
4993a40ed3dSBarry Smith   PetscFunctionBegin;
5006f08ad05SPierre Jolivet   if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
5019566063dSJacob Faibussowitsch   PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
502d016bddeSToby Isaac   PetscCall(VecGetArrayWriteAndMemType(mdn->lvec, &ay, &aymtype));
5039566063dSJacob Faibussowitsch   PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
5049566063dSJacob Faibussowitsch   PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
505d016bddeSToby Isaac   PetscCall(VecRestoreArrayWriteAndMemType(mdn->lvec, &ay));
5069566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
5079566063dSJacob Faibussowitsch   PetscCall((*mdn->A->ops->mult)(mdn->A, mdn->lvec, yy));
5083ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
5098965ea79SLois Curfman McInnes }
5108965ea79SLois Curfman McInnes 
5110be0d8bdSHansol Suh static PetscErrorCode MatMultAddColumnRange_MPIDense(Mat mat, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end)
5120be0d8bdSHansol Suh {
5130be0d8bdSHansol Suh   Mat_MPIDense      *mdn = (Mat_MPIDense *)mat->data;
5140be0d8bdSHansol Suh   const PetscScalar *ax;
5150be0d8bdSHansol Suh   PetscScalar       *ay;
5160be0d8bdSHansol Suh   PetscMemType       axmtype, aymtype;
5170be0d8bdSHansol Suh 
5180be0d8bdSHansol Suh   PetscFunctionBegin;
5190be0d8bdSHansol Suh   if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
5200be0d8bdSHansol Suh   PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
5210be0d8bdSHansol Suh   PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype));
5220be0d8bdSHansol Suh   PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
5230be0d8bdSHansol Suh   PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
5240be0d8bdSHansol Suh   PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay));
5250be0d8bdSHansol Suh   PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
5260be0d8bdSHansol Suh   PetscUseMethod(mdn->A, "MatMultAddColumnRange_C", (Mat, Vec, Vec, Vec, PetscInt, PetscInt), (mdn->A, mdn->lvec, yy, zz, c_start, c_end));
5270be0d8bdSHansol Suh   PetscFunctionReturn(PETSC_SUCCESS);
5280be0d8bdSHansol Suh }
5290be0d8bdSHansol Suh 
53066976f2fSJacob Faibussowitsch static PetscErrorCode MatMultAdd_MPIDense(Mat mat, Vec xx, Vec yy, Vec zz)
531d71ae5a4SJacob Faibussowitsch {
53239ddd567SLois Curfman McInnes   Mat_MPIDense      *mdn = (Mat_MPIDense *)mat->data;
533637a0070SStefano Zampini   const PetscScalar *ax;
534637a0070SStefano Zampini   PetscScalar       *ay;
535d0295fc0SJunchao Zhang   PetscMemType       axmtype, aymtype;
536c456f294SBarry Smith 
5373a40ed3dSBarry Smith   PetscFunctionBegin;
5386f08ad05SPierre Jolivet   if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
5399566063dSJacob Faibussowitsch   PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
5409566063dSJacob Faibussowitsch   PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype));
5419566063dSJacob Faibussowitsch   PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
5429566063dSJacob Faibussowitsch   PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
5439566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay));
5449566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
5459566063dSJacob Faibussowitsch   PetscCall((*mdn->A->ops->multadd)(mdn->A, mdn->lvec, yy, zz));
5463ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
5478965ea79SLois Curfman McInnes }
5488965ea79SLois Curfman McInnes 
5490be0d8bdSHansol Suh static PetscErrorCode MatMultHermitianTransposeColumnRange_MPIDense(Mat A, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end)
5500be0d8bdSHansol Suh {
5510be0d8bdSHansol Suh   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
5520be0d8bdSHansol Suh   const PetscScalar *ax;
5530be0d8bdSHansol Suh   PetscScalar       *ay;
5540be0d8bdSHansol Suh   PetscMemType       axmtype, aymtype;
555d016bddeSToby Isaac   PetscInt           r_start, r_end;
556d016bddeSToby Isaac   PetscInt           c_start_local, c_end_local;
5570be0d8bdSHansol Suh 
5580be0d8bdSHansol Suh   PetscFunctionBegin;
5590be0d8bdSHansol Suh   if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
560d016bddeSToby Isaac   PetscCall(VecZeroEntries(a->lvec));
561d016bddeSToby Isaac   PetscCall(VecGetOwnershipRange(yy, &r_start, &r_end));
562d016bddeSToby Isaac   c_start_local = PetscMax(c_start, r_start);
563d016bddeSToby Isaac   c_end_local   = PetscMin(c_end, r_end);
564d016bddeSToby Isaac   PetscCall(VecGetArrayAndMemType(yy, &ay, &aymtype));
565d016bddeSToby Isaac   if (c_end_local > c_start_local) {
566d016bddeSToby Isaac     if (PetscMemTypeHost(aymtype)) {
567d016bddeSToby Isaac       PetscCall(PetscArrayzero(&ay[c_start_local], (size_t)(c_end_local - c_start_local)));
568d016bddeSToby Isaac     } else {
569d016bddeSToby Isaac       PetscCall(PetscDeviceRegisterMemory(ay, aymtype, sizeof(*ay) * ((size_t)(r_end - r_start))));
570d016bddeSToby Isaac       PetscCall(PetscDeviceArrayZero(NULL, &ay[c_start_local], (size_t)(c_end_local - c_start_local)));
571d016bddeSToby Isaac     }
572d016bddeSToby Isaac   }
5730be0d8bdSHansol Suh   PetscUseMethod(a->A, "MatMultHermitianTransposeColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (a->A, xx, a->lvec, c_start, c_end));
5740be0d8bdSHansol Suh   PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
5750be0d8bdSHansol Suh   PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
5760be0d8bdSHansol Suh   PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
5770be0d8bdSHansol Suh   PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
5780be0d8bdSHansol Suh   PetscCall(VecRestoreArrayAndMemType(yy, &ay));
5790be0d8bdSHansol Suh   PetscFunctionReturn(PETSC_SUCCESS);
5800be0d8bdSHansol Suh }
5810be0d8bdSHansol Suh 
582459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTransposeKernel_MPIDense(Mat A, Vec xx, Vec yy, PetscBool herm)
583d71ae5a4SJacob Faibussowitsch {
584096963f5SLois Curfman McInnes   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
585637a0070SStefano Zampini   const PetscScalar *ax;
586637a0070SStefano Zampini   PetscScalar       *ay;
587d0295fc0SJunchao Zhang   PetscMemType       axmtype, aymtype;
5884d86920dSPierre Jolivet 
5893a40ed3dSBarry Smith   PetscFunctionBegin;
5906f08ad05SPierre Jolivet   if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
5919566063dSJacob Faibussowitsch   PetscCall(VecSet(yy, 0.0));
592459e8d23SBlanca Mellado Pinto   if (herm) PetscCall((*a->A->ops->multhermitiantranspose)(a->A, xx, a->lvec));
593459e8d23SBlanca Mellado Pinto   else PetscCall((*a->A->ops->multtranspose)(a->A, xx, a->lvec));
5949566063dSJacob Faibussowitsch   PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
5959566063dSJacob Faibussowitsch   PetscCall(VecGetArrayAndMemType(yy, &ay, &aymtype));
5969566063dSJacob Faibussowitsch   PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
5979566063dSJacob Faibussowitsch   PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
5989566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
5999566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayAndMemType(yy, &ay));
6003ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
601096963f5SLois Curfman McInnes }
602096963f5SLois Curfman McInnes 
6030be0d8bdSHansol Suh static PetscErrorCode MatMultHermitianTransposeAddColumnRange_MPIDense(Mat A, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end)
6040be0d8bdSHansol Suh {
6050be0d8bdSHansol Suh   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
6060be0d8bdSHansol Suh   const PetscScalar *ax;
6070be0d8bdSHansol Suh   PetscScalar       *ay;
6080be0d8bdSHansol Suh   PetscMemType       axmtype, aymtype;
6090be0d8bdSHansol Suh 
6100be0d8bdSHansol Suh   PetscFunctionBegin;
6110be0d8bdSHansol Suh   if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
6120be0d8bdSHansol Suh   PetscCall(VecCopy(yy, zz));
613d016bddeSToby Isaac   PetscCall(VecZeroEntries(a->lvec));
6140be0d8bdSHansol Suh   PetscUseMethod(a->A, "MatMultHermitianTransposeColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (a->A, xx, a->lvec, c_start, c_end));
6150be0d8bdSHansol Suh   PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
6160be0d8bdSHansol Suh   PetscCall(VecGetArrayAndMemType(zz, &ay, &aymtype));
6170be0d8bdSHansol Suh   PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
6180be0d8bdSHansol Suh   PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
6190be0d8bdSHansol Suh   PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
6200be0d8bdSHansol Suh   PetscCall(VecRestoreArrayAndMemType(zz, &ay));
6210be0d8bdSHansol Suh   PetscFunctionReturn(PETSC_SUCCESS);
6220be0d8bdSHansol Suh }
6230be0d8bdSHansol Suh 
624459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTransposeAddKernel_MPIDense(Mat A, Vec xx, Vec yy, Vec zz, PetscBool herm)
625d71ae5a4SJacob Faibussowitsch {
626096963f5SLois Curfman McInnes   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
627637a0070SStefano Zampini   const PetscScalar *ax;
628637a0070SStefano Zampini   PetscScalar       *ay;
629d0295fc0SJunchao Zhang   PetscMemType       axmtype, aymtype;
630096963f5SLois Curfman McInnes 
6313a40ed3dSBarry Smith   PetscFunctionBegin;
6326f08ad05SPierre Jolivet   if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
6339566063dSJacob Faibussowitsch   PetscCall(VecCopy(yy, zz));
634459e8d23SBlanca Mellado Pinto   if (herm) PetscCall((*a->A->ops->multhermitiantranspose)(a->A, xx, a->lvec));
635459e8d23SBlanca Mellado Pinto   else PetscCall((*a->A->ops->multtranspose)(a->A, xx, a->lvec));
6369566063dSJacob Faibussowitsch   PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
6379566063dSJacob Faibussowitsch   PetscCall(VecGetArrayAndMemType(zz, &ay, &aymtype));
6389566063dSJacob Faibussowitsch   PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
6399566063dSJacob Faibussowitsch   PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
6409566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
6419566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayAndMemType(zz, &ay));
6423ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
643096963f5SLois Curfman McInnes }
644096963f5SLois Curfman McInnes 
645459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTranspose_MPIDense(Mat A, Vec xx, Vec yy)
646459e8d23SBlanca Mellado Pinto {
647459e8d23SBlanca Mellado Pinto   PetscFunctionBegin;
648459e8d23SBlanca Mellado Pinto   PetscCall(MatMultTransposeKernel_MPIDense(A, xx, yy, PETSC_FALSE));
649459e8d23SBlanca Mellado Pinto   PetscFunctionReturn(PETSC_SUCCESS);
650459e8d23SBlanca Mellado Pinto }
651459e8d23SBlanca Mellado Pinto 
652459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTransposeAdd_MPIDense(Mat A, Vec xx, Vec yy, Vec zz)
653459e8d23SBlanca Mellado Pinto {
654459e8d23SBlanca Mellado Pinto   PetscFunctionBegin;
655459e8d23SBlanca Mellado Pinto   PetscCall(MatMultTransposeAddKernel_MPIDense(A, xx, yy, zz, PETSC_FALSE));
656459e8d23SBlanca Mellado Pinto   PetscFunctionReturn(PETSC_SUCCESS);
657459e8d23SBlanca Mellado Pinto }
658459e8d23SBlanca Mellado Pinto 
659459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultHermitianTranspose_MPIDense(Mat A, Vec xx, Vec yy)
660459e8d23SBlanca Mellado Pinto {
661459e8d23SBlanca Mellado Pinto   PetscFunctionBegin;
662459e8d23SBlanca Mellado Pinto   PetscCall(MatMultTransposeKernel_MPIDense(A, xx, yy, PETSC_TRUE));
663459e8d23SBlanca Mellado Pinto   PetscFunctionReturn(PETSC_SUCCESS);
664459e8d23SBlanca Mellado Pinto }
665459e8d23SBlanca Mellado Pinto 
666459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultHermitianTransposeAdd_MPIDense(Mat A, Vec xx, Vec yy, Vec zz)
667459e8d23SBlanca Mellado Pinto {
668459e8d23SBlanca Mellado Pinto   PetscFunctionBegin;
669459e8d23SBlanca Mellado Pinto   PetscCall(MatMultTransposeAddKernel_MPIDense(A, xx, yy, zz, PETSC_TRUE));
670459e8d23SBlanca Mellado Pinto   PetscFunctionReturn(PETSC_SUCCESS);
671459e8d23SBlanca Mellado Pinto }
672459e8d23SBlanca Mellado Pinto 
673d71ae5a4SJacob Faibussowitsch PetscErrorCode MatGetDiagonal_MPIDense(Mat A, Vec v)
674d71ae5a4SJacob Faibussowitsch {
67539ddd567SLois Curfman McInnes   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
67679560b15SJacob Faibussowitsch   PetscInt           lda, len, i, nl, ng, m = A->rmap->n, radd;
67779560b15SJacob Faibussowitsch   PetscScalar       *x;
678637a0070SStefano Zampini   const PetscScalar *av;
679ed3cc1f0SBarry Smith 
6803a40ed3dSBarry Smith   PetscFunctionBegin;
6819566063dSJacob Faibussowitsch   PetscCall(VecGetArray(v, &x));
68279560b15SJacob Faibussowitsch   PetscCall(VecGetSize(v, &ng));
68379560b15SJacob Faibussowitsch   PetscCheck(ng == A->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Nonconforming mat and vec");
68479560b15SJacob Faibussowitsch   PetscCall(VecGetLocalSize(v, &nl));
685d0f46423SBarry Smith   len  = PetscMin(a->A->rmap->n, a->A->cmap->n);
686d0f46423SBarry Smith   radd = A->rmap->rstart * m;
6879566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, &av));
6889566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
689ad540459SPierre Jolivet   for (i = 0; i < len; i++) x[i] = av[radd + i * lda + i];
6909566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, &av));
6910be0d8bdSHansol Suh   if (nl - i > 0) PetscCall(PetscArrayzero(x + i, nl - i));
6929566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(v, &x));
6933ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
6948965ea79SLois Curfman McInnes }
6958965ea79SLois Curfman McInnes 
69666976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MPIDense(Mat mat)
697d71ae5a4SJacob Faibussowitsch {
6983501a2bdSLois Curfman McInnes   Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
699ed3cc1f0SBarry Smith 
7003a40ed3dSBarry Smith   PetscFunctionBegin;
7013ba16761SJacob Faibussowitsch   PetscCall(PetscLogObjectState((PetscObject)mat, "Rows=%" PetscInt_FMT ", Cols=%" PetscInt_FMT, mat->rmap->N, mat->cmap->N));
7029566063dSJacob Faibussowitsch   PetscCall(MatStashDestroy_Private(&mat->stash));
70328b400f6SJacob Faibussowitsch   PetscCheck(!mdn->vecinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
70428b400f6SJacob Faibussowitsch   PetscCheck(!mdn->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
7059566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&mdn->A));
7069566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&mdn->lvec));
7079566063dSJacob Faibussowitsch   PetscCall(PetscSFDestroy(&mdn->Mvctx));
7089566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&mdn->cvec));
7099566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&mdn->cmat));
71001b82886SBarry Smith 
7119566063dSJacob Faibussowitsch   PetscCall(PetscFree(mat->data));
7129566063dSJacob Faibussowitsch   PetscCall(PetscObjectChangeTypeName((PetscObject)mat, NULL));
7138baccfbdSHong Zhang 
7149566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", NULL));
7159566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", NULL));
7169566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", NULL));
7179566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", NULL));
7189566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", NULL));
7199566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", NULL));
7209566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", NULL));
7219566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", NULL));
7229566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", NULL));
7239566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", NULL));
7249566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", NULL));
7259566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", NULL));
7269566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", NULL));
7278baccfbdSHong Zhang #if defined(PETSC_HAVE_ELEMENTAL)
7289566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", NULL));
7298baccfbdSHong Zhang #endif
730*d1a032dbSPierre Jolivet #if defined(PETSC_HAVE_SCALAPACK) && (defined(PETSC_USE_REAL_SINGLE) || defined(PETSC_USE_REAL_DOUBLE))
7319566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", NULL));
732d24d4204SJose E. Roman #endif
7339566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", NULL));
7349566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", NULL));
7359566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", NULL));
7366718818eSStefano Zampini #if defined(PETSC_HAVE_CUDA)
7379566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", NULL));
7389566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", NULL));
7392e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", NULL));
7402e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidensecuda_mpidense_C", NULL));
7412e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidensecuda_C", NULL));
7422e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidensecuda_C", NULL));
7432e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensecuda_mpiaij_C", NULL));
7442e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensecuda_mpiaijcusparse_C", NULL));
7452e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArray_C", NULL));
7462e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArrayRead_C", NULL));
7472e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArrayWrite_C", NULL));
7482e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArray_C", NULL));
7492e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArrayRead_C", NULL));
7502e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArrayWrite_C", NULL));
7512e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAPlaceArray_C", NULL));
7522e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAResetArray_C", NULL));
7532e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAReplaceArray_C", NULL));
7543d9668e3SJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDASetPreallocation_C", NULL));
7556718818eSStefano Zampini #endif
75647d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
75747d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", NULL));
75847d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", NULL));
75947d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", NULL));
76047d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidensehip_mpidense_C", NULL));
76147d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidensehip_C", NULL));
76247d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidensehip_C", NULL));
76347d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensehip_mpiaij_C", NULL));
76447d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensehip_mpiaijhipsparse_C", NULL));
76547d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArray_C", NULL));
76647d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArrayRead_C", NULL));
76747d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArrayWrite_C", NULL));
76847d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArray_C", NULL));
76947d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArrayRead_C", NULL));
77047d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArrayWrite_C", NULL));
77147d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPPlaceArray_C", NULL));
77247d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPResetArray_C", NULL));
77347d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPReplaceArray_C", NULL));
7743d9668e3SJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPSetPreallocation_C", NULL));
77547d993e7Ssuyashtn #endif
7769566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", NULL));
7779566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", NULL));
7789566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", NULL));
7799566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", NULL));
7809566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", NULL));
7819566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", NULL));
7829566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", NULL));
7839566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", NULL));
7849566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", NULL));
7859566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", NULL));
786d016bddeSToby Isaac   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultColumnRange_C", NULL));
7870be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", NULL));
7880be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", NULL));
7890be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", NULL));
790f8103e6bSStefano Zampini 
791f8103e6bSStefano Zampini   PetscCall(PetscObjectCompose((PetscObject)mat, "DiagonalBlock", NULL));
7923ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
7938965ea79SLois Curfman McInnes }
79439ddd567SLois Curfman McInnes 
7959804daf3SBarry Smith #include <petscdraw.h>
796d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatView_MPIDense_ASCIIorDraworSocket(Mat mat, PetscViewer viewer)
797d71ae5a4SJacob Faibussowitsch {
79839ddd567SLois Curfman McInnes   Mat_MPIDense     *mdn = (Mat_MPIDense *)mat->data;
799616b8fbbSStefano Zampini   PetscMPIInt       rank;
80019fd82e9SBarry Smith   PetscViewerType   vtype;
8019f196a02SMartin Diehl   PetscBool         isascii, isdraw;
802b0a32e0cSBarry Smith   PetscViewer       sviewer;
803f3ef73ceSBarry Smith   PetscViewerFormat format;
8048965ea79SLois Curfman McInnes 
8053a40ed3dSBarry Smith   PetscFunctionBegin;
8069566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mat), &rank));
8079f196a02SMartin Diehl   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
8089566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
8099f196a02SMartin Diehl   if (isascii) {
8109566063dSJacob Faibussowitsch     PetscCall(PetscViewerGetType(viewer, &vtype));
8119566063dSJacob Faibussowitsch     PetscCall(PetscViewerGetFormat(viewer, &format));
812456192e2SBarry Smith     if (format == PETSC_VIEWER_ASCII_INFO_DETAIL) {
8134e220ebcSLois Curfman McInnes       MatInfo info;
8149566063dSJacob Faibussowitsch       PetscCall(MatGetInfo(mat, MAT_LOCAL, &info));
8159566063dSJacob Faibussowitsch       PetscCall(PetscViewerASCIIPushSynchronized(viewer));
8169371c9d4SSatish 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,
8179371c9d4SSatish Balay                                                    (PetscInt)info.memory));
8189566063dSJacob Faibussowitsch       PetscCall(PetscViewerFlush(viewer));
8199566063dSJacob Faibussowitsch       PetscCall(PetscViewerASCIIPopSynchronized(viewer));
8206f08ad05SPierre Jolivet       if (mdn->Mvctx) PetscCall(PetscSFView(mdn->Mvctx, viewer));
8213ba16761SJacob Faibussowitsch       PetscFunctionReturn(PETSC_SUCCESS);
822fb9695e5SSatish Balay     } else if (format == PETSC_VIEWER_ASCII_INFO) {
8233ba16761SJacob Faibussowitsch       PetscFunctionReturn(PETSC_SUCCESS);
8248965ea79SLois Curfman McInnes     }
825f1af5d2fSBarry Smith   } else if (isdraw) {
826b0a32e0cSBarry Smith     PetscDraw draw;
827ace3abfcSBarry Smith     PetscBool isnull;
828f1af5d2fSBarry Smith 
8299566063dSJacob Faibussowitsch     PetscCall(PetscViewerDrawGetDraw(viewer, 0, &draw));
8309566063dSJacob Faibussowitsch     PetscCall(PetscDrawIsNull(draw, &isnull));
8313ba16761SJacob Faibussowitsch     if (isnull) PetscFunctionReturn(PETSC_SUCCESS);
832f1af5d2fSBarry Smith   }
83377ed5343SBarry Smith 
8347da1fb6eSBarry Smith   {
8358965ea79SLois Curfman McInnes     /* assemble the entire matrix onto first processor. */
8368965ea79SLois Curfman McInnes     Mat          A;
837d0f46423SBarry Smith     PetscInt     M = mat->rmap->N, N = mat->cmap->N, m, row, i, nz;
838ba8c8a56SBarry Smith     PetscInt    *cols;
839ba8c8a56SBarry Smith     PetscScalar *vals;
8408965ea79SLois Curfman McInnes 
8419566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)mat), &A));
842dd400576SPatrick Sanan     if (rank == 0) {
8439566063dSJacob Faibussowitsch       PetscCall(MatSetSizes(A, M, N, M, N));
8443a40ed3dSBarry Smith     } else {
8459566063dSJacob Faibussowitsch       PetscCall(MatSetSizes(A, 0, 0, M, N));
8468965ea79SLois Curfman McInnes     }
8477adad957SLisandro Dalcin     /* Since this is a temporary matrix, MATMPIDENSE instead of ((PetscObject)A)->type_name here is probably acceptable. */
8489566063dSJacob Faibussowitsch     PetscCall(MatSetType(A, MATMPIDENSE));
8499566063dSJacob Faibussowitsch     PetscCall(MatMPIDenseSetPreallocation(A, NULL));
8508965ea79SLois Curfman McInnes 
85139ddd567SLois Curfman McInnes     /* Copy the matrix ... This isn't the most efficient means,
85239ddd567SLois Curfman McInnes        but it's quick for now */
85351022da4SBarry Smith     A->insertmode = INSERT_VALUES;
8542205254eSKarl Rupp 
8552205254eSKarl Rupp     row = mat->rmap->rstart;
8562205254eSKarl Rupp     m   = mdn->A->rmap->n;
8578965ea79SLois Curfman McInnes     for (i = 0; i < m; i++) {
8589566063dSJacob Faibussowitsch       PetscCall(MatGetRow_MPIDense(mat, row, &nz, &cols, &vals));
8599566063dSJacob Faibussowitsch       PetscCall(MatSetValues_MPIDense(A, 1, &row, nz, cols, vals, INSERT_VALUES));
8609566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow_MPIDense(mat, row, &nz, &cols, &vals));
86139ddd567SLois Curfman McInnes       row++;
8628965ea79SLois Curfman McInnes     }
8638965ea79SLois Curfman McInnes 
8649566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
8659566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
8669566063dSJacob Faibussowitsch     PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
867dd400576SPatrick Sanan     if (rank == 0) {
868f4f49eeaSPierre Jolivet       PetscCall(PetscObjectSetName((PetscObject)((Mat_MPIDense *)A->data)->A, ((PetscObject)mat)->name));
869f4f49eeaSPierre Jolivet       PetscCall(MatView_SeqDense(((Mat_MPIDense *)A->data)->A, sviewer));
8708965ea79SLois Curfman McInnes     }
8719566063dSJacob Faibussowitsch     PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
8729566063dSJacob Faibussowitsch     PetscCall(MatDestroy(&A));
8738965ea79SLois Curfman McInnes   }
8743ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
8758965ea79SLois Curfman McInnes }
8768965ea79SLois Curfman McInnes 
87766976f2fSJacob Faibussowitsch static PetscErrorCode MatView_MPIDense(Mat mat, PetscViewer viewer)
878d71ae5a4SJacob Faibussowitsch {
8799f196a02SMartin Diehl   PetscBool isascii, isbinary, isdraw, issocket;
8808965ea79SLois Curfman McInnes 
881433994e6SBarry Smith   PetscFunctionBegin;
8829f196a02SMartin Diehl   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
8839566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
8849566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERSOCKET, &issocket));
8859566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
8860f5bd95cSBarry Smith 
887ac530a7eSPierre Jolivet   if (isascii || issocket || isdraw) PetscCall(MatView_MPIDense_ASCIIorDraworSocket(mat, viewer));
888ac530a7eSPierre Jolivet   else if (isbinary) PetscCall(MatView_Dense_Binary(mat, viewer));
8893ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
8908965ea79SLois Curfman McInnes }
8918965ea79SLois Curfman McInnes 
89266976f2fSJacob Faibussowitsch static PetscErrorCode MatGetInfo_MPIDense(Mat A, MatInfoType flag, MatInfo *info)
893d71ae5a4SJacob Faibussowitsch {
8943501a2bdSLois Curfman McInnes   Mat_MPIDense  *mat = (Mat_MPIDense *)A->data;
8953501a2bdSLois Curfman McInnes   Mat            mdn = mat->A;
8963966268fSBarry Smith   PetscLogDouble isend[5], irecv[5];
8978965ea79SLois Curfman McInnes 
8983a40ed3dSBarry Smith   PetscFunctionBegin;
8994e220ebcSLois Curfman McInnes   info->block_size = 1.0;
9002205254eSKarl Rupp 
9019566063dSJacob Faibussowitsch   PetscCall(MatGetInfo(mdn, MAT_LOCAL, info));
9022205254eSKarl Rupp 
9039371c9d4SSatish Balay   isend[0] = info->nz_used;
9049371c9d4SSatish Balay   isend[1] = info->nz_allocated;
9059371c9d4SSatish Balay   isend[2] = info->nz_unneeded;
9069371c9d4SSatish Balay   isend[3] = info->memory;
9079371c9d4SSatish Balay   isend[4] = info->mallocs;
9088965ea79SLois Curfman McInnes   if (flag == MAT_LOCAL) {
9094e220ebcSLois Curfman McInnes     info->nz_used      = isend[0];
9104e220ebcSLois Curfman McInnes     info->nz_allocated = isend[1];
9114e220ebcSLois Curfman McInnes     info->nz_unneeded  = isend[2];
9124e220ebcSLois Curfman McInnes     info->memory       = isend[3];
9134e220ebcSLois Curfman McInnes     info->mallocs      = isend[4];
9148965ea79SLois Curfman McInnes   } else if (flag == MAT_GLOBAL_MAX) {
915462c564dSBarry Smith     PetscCallMPI(MPIU_Allreduce(isend, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_MAX, PetscObjectComm((PetscObject)A)));
9162205254eSKarl Rupp 
9174e220ebcSLois Curfman McInnes     info->nz_used      = irecv[0];
9184e220ebcSLois Curfman McInnes     info->nz_allocated = irecv[1];
9194e220ebcSLois Curfman McInnes     info->nz_unneeded  = irecv[2];
9204e220ebcSLois Curfman McInnes     info->memory       = irecv[3];
9214e220ebcSLois Curfman McInnes     info->mallocs      = irecv[4];
9228965ea79SLois Curfman McInnes   } else if (flag == MAT_GLOBAL_SUM) {
923462c564dSBarry Smith     PetscCallMPI(MPIU_Allreduce(isend, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_SUM, PetscObjectComm((PetscObject)A)));
9242205254eSKarl Rupp 
9254e220ebcSLois Curfman McInnes     info->nz_used      = irecv[0];
9264e220ebcSLois Curfman McInnes     info->nz_allocated = irecv[1];
9274e220ebcSLois Curfman McInnes     info->nz_unneeded  = irecv[2];
9284e220ebcSLois Curfman McInnes     info->memory       = irecv[3];
9294e220ebcSLois Curfman McInnes     info->mallocs      = irecv[4];
9308965ea79SLois Curfman McInnes   }
9314e220ebcSLois Curfman McInnes   info->fill_ratio_given  = 0; /* no parallel LU/ILU/Cholesky */
9324e220ebcSLois Curfman McInnes   info->fill_ratio_needed = 0;
9334e220ebcSLois Curfman McInnes   info->factor_mallocs    = 0;
9343ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9358965ea79SLois Curfman McInnes }
9368965ea79SLois Curfman McInnes 
93766976f2fSJacob Faibussowitsch static PetscErrorCode MatSetOption_MPIDense(Mat A, MatOption op, PetscBool flg)
938d71ae5a4SJacob Faibussowitsch {
93939ddd567SLois Curfman McInnes   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
9408965ea79SLois Curfman McInnes 
9413a40ed3dSBarry Smith   PetscFunctionBegin;
94212c028f9SKris Buschelman   switch (op) {
943512a5fc5SBarry Smith   case MAT_NEW_NONZERO_LOCATIONS:
94412c028f9SKris Buschelman   case MAT_NEW_NONZERO_LOCATION_ERR:
94512c028f9SKris Buschelman   case MAT_NEW_NONZERO_ALLOCATION_ERR:
94643674050SBarry Smith     MatCheckPreallocated(A, 1);
9479566063dSJacob Faibussowitsch     PetscCall(MatSetOption(a->A, op, flg));
94812c028f9SKris Buschelman     break;
94912c028f9SKris Buschelman   case MAT_ROW_ORIENTED:
95043674050SBarry Smith     MatCheckPreallocated(A, 1);
9514e0d8c25SBarry Smith     a->roworiented = flg;
9529566063dSJacob Faibussowitsch     PetscCall(MatSetOption(a->A, op, flg));
95312c028f9SKris Buschelman     break;
954d71ae5a4SJacob Faibussowitsch   case MAT_IGNORE_OFF_PROC_ENTRIES:
955d71ae5a4SJacob Faibussowitsch     a->donotstash = flg;
956d71ae5a4SJacob Faibussowitsch     break;
95777e54ba9SKris Buschelman   case MAT_SYMMETRIC:
95877e54ba9SKris Buschelman   case MAT_STRUCTURALLY_SYMMETRIC:
9599a4540c5SBarry Smith   case MAT_HERMITIAN:
9609a4540c5SBarry Smith   case MAT_SYMMETRY_ETERNAL:
961b94d7dedSBarry Smith   case MAT_STRUCTURAL_SYMMETRY_ETERNAL:
962b94d7dedSBarry Smith   case MAT_SPD:
963b94d7dedSBarry Smith   case MAT_SPD_ETERNAL:
964b94d7dedSBarry Smith     /* if the diagonal matrix is square it inherits some of the properties above */
9659725ddffSPierre Jolivet     if (a->A && A->rmap->n == A->cmap->n) PetscCall(MatSetOption(a->A, op, flg));
96677e54ba9SKris Buschelman     break;
967d71ae5a4SJacob Faibussowitsch   default:
968888c827cSStefano Zampini     break;
9693a40ed3dSBarry Smith   }
9703ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9718965ea79SLois Curfman McInnes }
9728965ea79SLois Curfman McInnes 
97366976f2fSJacob Faibussowitsch static PetscErrorCode MatDiagonalScale_MPIDense(Mat A, Vec ll, Vec rr)
974d71ae5a4SJacob Faibussowitsch {
9755b2fa520SLois Curfman McInnes   Mat_MPIDense      *mdn = (Mat_MPIDense *)A->data;
976637a0070SStefano Zampini   const PetscScalar *l;
977637a0070SStefano Zampini   PetscScalar        x, *v, *vv, *r;
978637a0070SStefano Zampini   PetscInt           i, j, s2a, s3a, s2, s3, m = mdn->A->rmap->n, n = mdn->A->cmap->n, lda;
9795b2fa520SLois Curfman McInnes 
9805b2fa520SLois Curfman McInnes   PetscFunctionBegin;
9819566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(mdn->A, &vv));
9829566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(mdn->A, &lda));
9839566063dSJacob Faibussowitsch   PetscCall(MatGetLocalSize(A, &s2, &s3));
9845b2fa520SLois Curfman McInnes   if (ll) {
9859566063dSJacob Faibussowitsch     PetscCall(VecGetLocalSize(ll, &s2a));
98608401ef6SPierre Jolivet     PetscCheck(s2a == s2, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Left scaling vector non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT, s2a, s2);
9879566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(ll, &l));
9885b2fa520SLois Curfman McInnes     for (i = 0; i < m; i++) {
9895b2fa520SLois Curfman McInnes       x = l[i];
990637a0070SStefano Zampini       v = vv + i;
9919371c9d4SSatish Balay       for (j = 0; j < n; j++) {
9929371c9d4SSatish Balay         (*v) *= x;
9939371c9d4SSatish Balay         v += lda;
9949371c9d4SSatish Balay       }
9955b2fa520SLois Curfman McInnes     }
9969566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(ll, &l));
9979566063dSJacob Faibussowitsch     PetscCall(PetscLogFlops(1.0 * n * m));
9985b2fa520SLois Curfman McInnes   }
9995b2fa520SLois Curfman McInnes   if (rr) {
1000637a0070SStefano Zampini     const PetscScalar *ar;
1001637a0070SStefano Zampini 
10029566063dSJacob Faibussowitsch     PetscCall(VecGetLocalSize(rr, &s3a));
100308401ef6SPierre Jolivet     PetscCheck(s3a == s3, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Right scaling vec non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT ".", s3a, s3);
10049566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(rr, &ar));
10056f08ad05SPierre Jolivet     if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
10069566063dSJacob Faibussowitsch     PetscCall(VecGetArray(mdn->lvec, &r));
10079566063dSJacob Faibussowitsch     PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, ar, r, MPI_REPLACE));
10089566063dSJacob Faibussowitsch     PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ar, r, MPI_REPLACE));
10099566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(rr, &ar));
10105b2fa520SLois Curfman McInnes     for (i = 0; i < n; i++) {
10115b2fa520SLois Curfman McInnes       x = r[i];
1012637a0070SStefano Zampini       v = vv + i * lda;
10132205254eSKarl Rupp       for (j = 0; j < m; j++) (*v++) *= x;
10145b2fa520SLois Curfman McInnes     }
10159566063dSJacob Faibussowitsch     PetscCall(VecRestoreArray(mdn->lvec, &r));
10169566063dSJacob Faibussowitsch     PetscCall(PetscLogFlops(1.0 * n * m));
10175b2fa520SLois Curfman McInnes   }
10189566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(mdn->A, &vv));
10193ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
10205b2fa520SLois Curfman McInnes }
10215b2fa520SLois Curfman McInnes 
102266976f2fSJacob Faibussowitsch static PetscErrorCode MatNorm_MPIDense(Mat A, NormType type, PetscReal *nrm)
1023d71ae5a4SJacob Faibussowitsch {
10243501a2bdSLois Curfman McInnes   Mat_MPIDense      *mdn = (Mat_MPIDense *)A->data;
102513f74950SBarry Smith   PetscInt           i, j;
1026e91c04dfSPierre Jolivet   PetscMPIInt        size;
1027329f5518SBarry Smith   PetscReal          sum = 0.0;
1028637a0070SStefano Zampini   const PetscScalar *av, *v;
10293501a2bdSLois Curfman McInnes 
10303a40ed3dSBarry Smith   PetscFunctionBegin;
10319566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(mdn->A, &av));
1032637a0070SStefano Zampini   v = av;
10339566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
1034616b8fbbSStefano Zampini   if (size == 1) {
10359566063dSJacob Faibussowitsch     PetscCall(MatNorm(mdn->A, type, nrm));
10363501a2bdSLois Curfman McInnes   } else {
10373501a2bdSLois Curfman McInnes     if (type == NORM_FROBENIUS) {
1038d0f46423SBarry Smith       for (i = 0; i < mdn->A->cmap->n * mdn->A->rmap->n; i++) {
10399371c9d4SSatish Balay         sum += PetscRealPart(PetscConj(*v) * (*v));
10409371c9d4SSatish Balay         v++;
10413501a2bdSLois Curfman McInnes       }
1042462c564dSBarry Smith       PetscCallMPI(MPIU_Allreduce(&sum, nrm, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
10438f1a2a5eSBarry Smith       *nrm = PetscSqrtReal(*nrm);
10449566063dSJacob Faibussowitsch       PetscCall(PetscLogFlops(2.0 * mdn->A->cmap->n * mdn->A->rmap->n));
10453a40ed3dSBarry Smith     } else if (type == NORM_1) {
1046e91c04dfSPierre Jolivet       PetscReal *tmp;
1047e91c04dfSPierre Jolivet 
1048e91c04dfSPierre Jolivet       PetscCall(PetscCalloc1(A->cmap->N, &tmp));
1049064f8208SBarry Smith       *nrm = 0.0;
1050637a0070SStefano Zampini       v    = av;
1051d0f46423SBarry Smith       for (j = 0; j < mdn->A->cmap->n; j++) {
1052d0f46423SBarry Smith         for (i = 0; i < mdn->A->rmap->n; i++) {
10539371c9d4SSatish Balay           tmp[j] += PetscAbsScalar(*v);
10549371c9d4SSatish Balay           v++;
10553501a2bdSLois Curfman McInnes         }
10563501a2bdSLois Curfman McInnes       }
1057e91c04dfSPierre Jolivet       PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, tmp, A->cmap->N, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
1058d0f46423SBarry Smith       for (j = 0; j < A->cmap->N; j++) {
1059e91c04dfSPierre Jolivet         if (tmp[j] > *nrm) *nrm = tmp[j];
10603501a2bdSLois Curfman McInnes       }
1061e91c04dfSPierre Jolivet       PetscCall(PetscFree(tmp));
10629566063dSJacob Faibussowitsch       PetscCall(PetscLogFlops(A->cmap->n * A->rmap->n));
10633a40ed3dSBarry Smith     } else if (type == NORM_INFINITY) { /* max row norm */
1064e91c04dfSPierre Jolivet       PetscCall(MatNorm(mdn->A, type, nrm));
1065e91c04dfSPierre Jolivet       PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, nrm, 1, MPIU_REAL, MPIU_MAX, PetscObjectComm((PetscObject)A)));
1066ce94432eSBarry Smith     } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "No support for two norm");
10673501a2bdSLois Curfman McInnes   }
10689566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(mdn->A, &av));
10693ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
10703501a2bdSLois Curfman McInnes }
10713501a2bdSLois Curfman McInnes 
107266976f2fSJacob Faibussowitsch static PetscErrorCode MatTranspose_MPIDense(Mat A, MatReuse reuse, Mat *matout)
1073d71ae5a4SJacob Faibussowitsch {
10743501a2bdSLois Curfman McInnes   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
10753501a2bdSLois Curfman McInnes   Mat           B;
1076d0f46423SBarry Smith   PetscInt      M = A->rmap->N, N = A->cmap->N, m, n, *rwork, rstart = A->rmap->rstart;
1077637a0070SStefano Zampini   PetscInt      j, i, lda;
107887828ca2SBarry Smith   PetscScalar  *v;
10793501a2bdSLois Curfman McInnes 
10803a40ed3dSBarry Smith   PetscFunctionBegin;
10817fb60732SBarry Smith   if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(A, *matout));
1082cf37664fSBarry Smith   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX) {
10839566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
10849566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(B, A->cmap->n, A->rmap->n, N, M));
10859566063dSJacob Faibussowitsch     PetscCall(MatSetType(B, ((PetscObject)A)->type_name));
10869566063dSJacob Faibussowitsch     PetscCall(MatMPIDenseSetPreallocation(B, NULL));
1087637a0070SStefano Zampini   } else B = *matout;
10883501a2bdSLois Curfman McInnes 
10899371c9d4SSatish Balay   m = a->A->rmap->n;
10909371c9d4SSatish Balay   n = a->A->cmap->n;
10919566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)&v));
10929566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
10939566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(m, &rwork));
10943501a2bdSLois Curfman McInnes   for (i = 0; i < m; i++) rwork[i] = rstart + i;
10951acff37aSSatish Balay   for (j = 0; j < n; j++) {
10969566063dSJacob Faibussowitsch     PetscCall(MatSetValues(B, 1, &j, m, rwork, v, INSERT_VALUES));
10978e3a54c0SPierre Jolivet     v = PetscSafePointerPlusOffset(v, lda);
10983501a2bdSLois Curfman McInnes   }
10999566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)&v));
11009566063dSJacob Faibussowitsch   PetscCall(PetscFree(rwork));
11019566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
11029566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
1103cf37664fSBarry Smith   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_REUSE_MATRIX) {
11043501a2bdSLois Curfman McInnes     *matout = B;
11053501a2bdSLois Curfman McInnes   } else {
11069566063dSJacob Faibussowitsch     PetscCall(MatHeaderMerge(A, &B));
11073501a2bdSLois Curfman McInnes   }
11083ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1109096963f5SLois Curfman McInnes }
1110096963f5SLois Curfman McInnes 
11116849ba73SBarry Smith static PetscErrorCode       MatDuplicate_MPIDense(Mat, MatDuplicateOption, Mat *);
111252c5f739Sprj- PETSC_INTERN PetscErrorCode MatScale_MPIDense(Mat, PetscScalar);
11138965ea79SLois Curfman McInnes 
111466976f2fSJacob Faibussowitsch static PetscErrorCode MatSetUp_MPIDense(Mat A)
1115d71ae5a4SJacob Faibussowitsch {
1116273d9f13SBarry Smith   PetscFunctionBegin;
11179566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(A->rmap));
11189566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(A->cmap));
111948a46eb9SPierre Jolivet   if (!A->preallocated) PetscCall(MatMPIDenseSetPreallocation(A, NULL));
11203ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1121273d9f13SBarry Smith }
1122273d9f13SBarry Smith 
112366976f2fSJacob Faibussowitsch static PetscErrorCode MatAXPY_MPIDense(Mat Y, PetscScalar alpha, Mat X, MatStructure str)
1124d71ae5a4SJacob Faibussowitsch {
1125488007eeSBarry Smith   Mat_MPIDense *A = (Mat_MPIDense *)Y->data, *B = (Mat_MPIDense *)X->data;
1126488007eeSBarry Smith 
1127488007eeSBarry Smith   PetscFunctionBegin;
11289566063dSJacob Faibussowitsch   PetscCall(MatAXPY(A->A, alpha, B->A, str));
11293ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1130488007eeSBarry Smith }
1131488007eeSBarry Smith 
113266976f2fSJacob Faibussowitsch static PetscErrorCode MatConjugate_MPIDense(Mat mat)
1133d71ae5a4SJacob Faibussowitsch {
1134ba337c44SJed Brown   Mat_MPIDense *a = (Mat_MPIDense *)mat->data;
1135ba337c44SJed Brown 
1136ba337c44SJed Brown   PetscFunctionBegin;
11379566063dSJacob Faibussowitsch   PetscCall(MatConjugate(a->A));
11383ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1139ba337c44SJed Brown }
1140ba337c44SJed Brown 
114166976f2fSJacob Faibussowitsch static PetscErrorCode MatRealPart_MPIDense(Mat A)
1142d71ae5a4SJacob Faibussowitsch {
1143ba337c44SJed Brown   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1144ba337c44SJed Brown 
1145ba337c44SJed Brown   PetscFunctionBegin;
11469566063dSJacob Faibussowitsch   PetscCall(MatRealPart(a->A));
11473ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1148ba337c44SJed Brown }
1149ba337c44SJed Brown 
115066976f2fSJacob Faibussowitsch static PetscErrorCode MatImaginaryPart_MPIDense(Mat A)
1151d71ae5a4SJacob Faibussowitsch {
1152ba337c44SJed Brown   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1153ba337c44SJed Brown 
1154ba337c44SJed Brown   PetscFunctionBegin;
11559566063dSJacob Faibussowitsch   PetscCall(MatImaginaryPart(a->A));
11563ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1157ba337c44SJed Brown }
1158ba337c44SJed Brown 
1159d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatGetColumnVector_MPIDense(Mat A, Vec v, PetscInt col)
1160d71ae5a4SJacob Faibussowitsch {
1161637a0070SStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
116249a6ff4bSBarry Smith 
116349a6ff4bSBarry Smith   PetscFunctionBegin;
116428b400f6SJacob Faibussowitsch   PetscCheck(a->A, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing local matrix");
116528b400f6SJacob Faibussowitsch   PetscCheck(a->A->ops->getcolumnvector, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing get column operation");
11669566063dSJacob Faibussowitsch   PetscCall((*a->A->ops->getcolumnvector)(a->A, v, col));
11673ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
116849a6ff4bSBarry Smith }
116949a6ff4bSBarry Smith 
1170857cbf51SRichard Tran Mills PETSC_INTERN PetscErrorCode MatGetColumnReductions_SeqDense(Mat, PetscInt, PetscReal *);
117152c5f739Sprj- 
117266976f2fSJacob Faibussowitsch static PetscErrorCode MatGetColumnReductions_MPIDense(Mat A, PetscInt type, PetscReal *reductions)
1173d71ae5a4SJacob Faibussowitsch {
1174a873a8cdSSam Reynolds   PetscInt      i, m, n;
11750716a85fSBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
11760716a85fSBarry Smith 
11770716a85fSBarry Smith   PetscFunctionBegin;
11789566063dSJacob Faibussowitsch   PetscCall(MatGetSize(A, &m, &n));
1179857cbf51SRichard Tran Mills   if (type == REDUCTION_MEAN_REALPART) {
1180e91c04dfSPierre Jolivet     PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_REALPART, reductions));
1181857cbf51SRichard Tran Mills   } else if (type == REDUCTION_MEAN_IMAGINARYPART) {
1182e91c04dfSPierre Jolivet     PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_IMAGINARYPART, reductions));
1183a873a8cdSSam Reynolds   } else {
1184e91c04dfSPierre Jolivet     PetscCall(MatGetColumnReductions_SeqDense(a->A, type, reductions));
1185a873a8cdSSam Reynolds   }
1186857cbf51SRichard Tran Mills   if (type == NORM_2) {
1187e91c04dfSPierre Jolivet     for (i = 0; i < n; i++) reductions[i] *= reductions[i];
11880716a85fSBarry Smith   }
1189857cbf51SRichard Tran Mills   if (type == NORM_INFINITY) {
1190e91c04dfSPierre Jolivet     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, reductions, n, MPIU_REAL, MPIU_MAX, A->hdr.comm));
11910716a85fSBarry Smith   } else {
1192e91c04dfSPierre Jolivet     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, reductions, n, MPIU_REAL, MPIU_SUM, A->hdr.comm));
11930716a85fSBarry Smith   }
1194857cbf51SRichard Tran Mills   if (type == NORM_2) {
1195a873a8cdSSam Reynolds     for (i = 0; i < n; i++) reductions[i] = PetscSqrtReal(reductions[i]);
1196857cbf51SRichard Tran Mills   } else if (type == REDUCTION_MEAN_REALPART || type == REDUCTION_MEAN_IMAGINARYPART) {
1197a873a8cdSSam Reynolds     for (i = 0; i < n; i++) reductions[i] /= m;
11980716a85fSBarry Smith   }
11993ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
12000716a85fSBarry Smith }
12010716a85fSBarry Smith 
1202d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatSetRandom_MPIDense(Mat x, PetscRandom rctx)
1203d71ae5a4SJacob Faibussowitsch {
120473a71a0fSBarry Smith   Mat_MPIDense *d = (Mat_MPIDense *)x->data;
120573a71a0fSBarry Smith 
120673a71a0fSBarry Smith   PetscFunctionBegin;
12079566063dSJacob Faibussowitsch   PetscCall(MatSetRandom(d->A, rctx));
12083faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE)
12093faff063SStefano Zampini   x->offloadmask = d->A->offloadmask;
12103faff063SStefano Zampini #endif
12113ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
121273a71a0fSBarry Smith }
121373a71a0fSBarry Smith 
1214d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMissingDiagonal_MPIDense(Mat A, PetscBool *missing, PetscInt *d)
1215d71ae5a4SJacob Faibussowitsch {
12163b49f96aSBarry Smith   PetscFunctionBegin;
12173b49f96aSBarry Smith   *missing = PETSC_FALSE;
12183ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
12193b49f96aSBarry Smith }
12203b49f96aSBarry Smith 
12214222ddf1SHong Zhang static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat);
1222cc48ffa7SToby Isaac static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat);
12236718818eSStefano Zampini static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat);
1224a1f9a5e6SStefano Zampini static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat);
12256718818eSStefano Zampini static PetscErrorCode MatEqual_MPIDense(Mat, Mat, PetscBool *);
12266718818eSStefano Zampini static PetscErrorCode MatLoad_MPIDense(Mat, PetscViewer);
12275d8c7819SPierre Jolivet static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat);
1228cc48ffa7SToby Isaac 
122909dc0095SBarry Smith static struct _MatOps MatOps_Values = {MatSetValues_MPIDense,
123009dc0095SBarry Smith                                        MatGetRow_MPIDense,
123109dc0095SBarry Smith                                        MatRestoreRow_MPIDense,
123209dc0095SBarry Smith                                        MatMult_MPIDense,
123397304618SKris Buschelman                                        /*  4*/ MatMultAdd_MPIDense,
12347c922b88SBarry Smith                                        MatMultTranspose_MPIDense,
12357c922b88SBarry Smith                                        MatMultTransposeAdd_MPIDense,
1236f4259b30SLisandro Dalcin                                        NULL,
1237f4259b30SLisandro Dalcin                                        NULL,
1238f4259b30SLisandro Dalcin                                        NULL,
1239f4259b30SLisandro Dalcin                                        /* 10*/ NULL,
1240f4259b30SLisandro Dalcin                                        NULL,
1241f4259b30SLisandro Dalcin                                        NULL,
1242f4259b30SLisandro Dalcin                                        NULL,
124309dc0095SBarry Smith                                        MatTranspose_MPIDense,
124497304618SKris Buschelman                                        /* 15*/ MatGetInfo_MPIDense,
12456e4ee0c6SHong Zhang                                        MatEqual_MPIDense,
124609dc0095SBarry Smith                                        MatGetDiagonal_MPIDense,
12475b2fa520SLois Curfman McInnes                                        MatDiagonalScale_MPIDense,
124809dc0095SBarry Smith                                        MatNorm_MPIDense,
124997304618SKris Buschelman                                        /* 20*/ MatAssemblyBegin_MPIDense,
125009dc0095SBarry Smith                                        MatAssemblyEnd_MPIDense,
125109dc0095SBarry Smith                                        MatSetOption_MPIDense,
125209dc0095SBarry Smith                                        MatZeroEntries_MPIDense,
1253d519adbfSMatthew Knepley                                        /* 24*/ MatZeroRows_MPIDense,
1254f4259b30SLisandro Dalcin                                        NULL,
1255f4259b30SLisandro Dalcin                                        NULL,
1256f4259b30SLisandro Dalcin                                        NULL,
1257f4259b30SLisandro Dalcin                                        NULL,
12584994cf47SJed Brown                                        /* 29*/ MatSetUp_MPIDense,
1259f4259b30SLisandro Dalcin                                        NULL,
1260f4259b30SLisandro Dalcin                                        NULL,
1261c56a70eeSBarry Smith                                        MatGetDiagonalBlock_MPIDense,
1262f4259b30SLisandro Dalcin                                        NULL,
1263d519adbfSMatthew Knepley                                        /* 34*/ MatDuplicate_MPIDense,
1264f4259b30SLisandro Dalcin                                        NULL,
1265f4259b30SLisandro Dalcin                                        NULL,
1266f4259b30SLisandro Dalcin                                        NULL,
1267f4259b30SLisandro Dalcin                                        NULL,
1268d519adbfSMatthew Knepley                                        /* 39*/ MatAXPY_MPIDense,
12697dae84e0SHong Zhang                                        MatCreateSubMatrices_MPIDense,
1270f4259b30SLisandro Dalcin                                        NULL,
127109dc0095SBarry Smith                                        MatGetValues_MPIDense,
1272a76f77c3SStefano Zampini                                        MatCopy_MPIDense,
1273f4259b30SLisandro Dalcin                                        /* 44*/ NULL,
127409dc0095SBarry Smith                                        MatScale_MPIDense,
12752f605a99SJose E. Roman                                        MatShift_MPIDense,
1276f4259b30SLisandro Dalcin                                        NULL,
1277f4259b30SLisandro Dalcin                                        NULL,
127873a71a0fSBarry Smith                                        /* 49*/ MatSetRandom_MPIDense,
1279f4259b30SLisandro Dalcin                                        NULL,
1280f4259b30SLisandro Dalcin                                        NULL,
1281f4259b30SLisandro Dalcin                                        NULL,
1282f4259b30SLisandro Dalcin                                        NULL,
1283f4259b30SLisandro Dalcin                                        /* 54*/ NULL,
1284f4259b30SLisandro Dalcin                                        NULL,
1285f4259b30SLisandro Dalcin                                        NULL,
1286f4259b30SLisandro Dalcin                                        NULL,
1287f4259b30SLisandro Dalcin                                        NULL,
12887dae84e0SHong Zhang                                        /* 59*/ MatCreateSubMatrix_MPIDense,
1289b9b97703SBarry Smith                                        MatDestroy_MPIDense,
1290b9b97703SBarry Smith                                        MatView_MPIDense,
1291f4259b30SLisandro Dalcin                                        NULL,
1292f4259b30SLisandro Dalcin                                        NULL,
1293f4259b30SLisandro Dalcin                                        /* 64*/ NULL,
1294f4259b30SLisandro Dalcin                                        NULL,
1295f4259b30SLisandro Dalcin                                        NULL,
1296f4259b30SLisandro Dalcin                                        NULL,
1297f4259b30SLisandro Dalcin                                        NULL,
1298f4259b30SLisandro Dalcin                                        /* 69*/ NULL,
1299f4259b30SLisandro Dalcin                                        NULL,
1300f4259b30SLisandro Dalcin                                        NULL,
1301f4259b30SLisandro Dalcin                                        NULL,
1302f4259b30SLisandro Dalcin                                        NULL,
1303f4259b30SLisandro Dalcin                                        /* 74*/ NULL,
1304f4259b30SLisandro Dalcin                                        NULL,
1305f4259b30SLisandro Dalcin                                        NULL,
1306f4259b30SLisandro Dalcin                                        NULL,
13078bb0f5c6SPierre Jolivet                                        MatLoad_MPIDense,
1308f4259b30SLisandro Dalcin                                        /* 79*/ NULL,
1309f4259b30SLisandro Dalcin                                        NULL,
1310f4259b30SLisandro Dalcin                                        NULL,
1311f4259b30SLisandro Dalcin                                        NULL,
13128bb0f5c6SPierre Jolivet                                        /* 83*/ NULL,
1313f4259b30SLisandro Dalcin                                        NULL,
1314f4259b30SLisandro Dalcin                                        NULL,
1315f4259b30SLisandro Dalcin                                        NULL,
13166718818eSStefano Zampini                                        MatMatTransposeMultSymbolic_MPIDense_MPIDense,
1317cc48ffa7SToby Isaac                                        MatMatTransposeMultNumeric_MPIDense_MPIDense,
13188bb0f5c6SPierre Jolivet                                        /* 89*/ NULL,
13198bb0f5c6SPierre Jolivet                                        MatProductSetFromOptions_MPIDense,
1320f4259b30SLisandro Dalcin                                        NULL,
1321f4259b30SLisandro Dalcin                                        NULL,
1322ba337c44SJed Brown                                        MatConjugate_MPIDense,
13238bb0f5c6SPierre Jolivet                                        /* 94*/ NULL,
1324f4259b30SLisandro Dalcin                                        NULL,
1325ba337c44SJed Brown                                        MatRealPart_MPIDense,
1326ba337c44SJed Brown                                        MatImaginaryPart_MPIDense,
1327f4259b30SLisandro Dalcin                                        NULL,
13288bb0f5c6SPierre Jolivet                                        /*99*/ NULL,
13298bb0f5c6SPierre Jolivet                                        NULL,
13308bb0f5c6SPierre Jolivet                                        NULL,
13318bb0f5c6SPierre Jolivet                                        NULL,
13328bb0f5c6SPierre Jolivet                                        MatGetColumnVector_MPIDense,
13338bb0f5c6SPierre Jolivet                                        /*104*/ MatMissingDiagonal_MPIDense,
13348bb0f5c6SPierre Jolivet                                        NULL,
13358bb0f5c6SPierre Jolivet                                        NULL,
13368bb0f5c6SPierre Jolivet                                        NULL,
1337f4259b30SLisandro Dalcin                                        NULL,
1338f4259b30SLisandro Dalcin                                        /*109*/ NULL,
1339f4259b30SLisandro Dalcin                                        NULL,
1340f4259b30SLisandro Dalcin                                        NULL,
1341459e8d23SBlanca Mellado Pinto                                        MatMultHermitianTranspose_MPIDense,
1342459e8d23SBlanca Mellado Pinto                                        MatMultHermitianTransposeAdd_MPIDense,
13438bb0f5c6SPierre Jolivet                                        /*114*/ NULL,
13448bb0f5c6SPierre Jolivet                                        NULL,
13458bb0f5c6SPierre Jolivet                                        MatGetColumnReductions_MPIDense,
13468bb0f5c6SPierre Jolivet                                        NULL,
13478bb0f5c6SPierre Jolivet                                        NULL,
13488bb0f5c6SPierre Jolivet                                        NULL,
13498bb0f5c6SPierre Jolivet                                        /*120*/ NULL,
13508bb0f5c6SPierre Jolivet                                        MatTransposeMatMultSymbolic_MPIDense_MPIDense,
13518bb0f5c6SPierre Jolivet                                        MatTransposeMatMultNumeric_MPIDense_MPIDense,
1352f4259b30SLisandro Dalcin                                        NULL,
1353f4259b30SLisandro Dalcin                                        /*124*/ NULL,
13548bb0f5c6SPierre Jolivet                                        NULL,
1355f4259b30SLisandro Dalcin                                        NULL,
1356f4259b30SLisandro Dalcin                                        NULL,
1357f4259b30SLisandro Dalcin                                        NULL,
1358f4259b30SLisandro Dalcin                                        /*129*/ NULL,
1359f4259b30SLisandro Dalcin                                        NULL,
13608bb0f5c6SPierre Jolivet                                        MatCreateMPIMatConcatenateSeqMat_MPIDense,
13618bb0f5c6SPierre Jolivet                                        NULL,
1362f4259b30SLisandro Dalcin                                        NULL,
1363f4259b30SLisandro Dalcin                                        /*134*/ NULL,
1364f4259b30SLisandro Dalcin                                        NULL,
1365f4259b30SLisandro Dalcin                                        NULL,
1366f4259b30SLisandro Dalcin                                        NULL,
1367f4259b30SLisandro Dalcin                                        NULL,
1368f4259b30SLisandro Dalcin                                        /*139*/ NULL,
1369f4259b30SLisandro Dalcin                                        NULL,
1370f4259b30SLisandro Dalcin                                        NULL,
137103db1824SAlex Lindsay                                        NULL,
1372c2be7ffeSStefano Zampini                                        NULL,
1373dec0b466SHong Zhang                                        NULL};
13748965ea79SLois Curfman McInnes 
137566976f2fSJacob Faibussowitsch static PetscErrorCode MatMPIDenseSetPreallocation_MPIDense(Mat mat, PetscScalar *data)
1376d71ae5a4SJacob Faibussowitsch {
1377637a0070SStefano Zampini   Mat_MPIDense *a     = (Mat_MPIDense *)mat->data;
137847d993e7Ssuyashtn   MatType       mtype = MATSEQDENSE;
1379a23d5eceSKris Buschelman 
1380a23d5eceSKris Buschelman   PetscFunctionBegin;
138128b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
13829566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(mat->rmap));
13839566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(mat->cmap));
1384637a0070SStefano Zampini   if (!a->A) {
13859566063dSJacob Faibussowitsch     PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
13869566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(a->A, mat->rmap->n, mat->cmap->N, mat->rmap->n, mat->cmap->N));
1387637a0070SStefano Zampini   }
138847d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA)
138947d993e7Ssuyashtn   PetscBool iscuda;
13909566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSECUDA, &iscuda));
139147d993e7Ssuyashtn   if (iscuda) mtype = MATSEQDENSECUDA;
139247d993e7Ssuyashtn #endif
139347d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
139447d993e7Ssuyashtn   PetscBool iship;
139547d993e7Ssuyashtn   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSEHIP, &iship));
139647d993e7Ssuyashtn   if (iship) mtype = MATSEQDENSEHIP;
139747d993e7Ssuyashtn #endif
139847d993e7Ssuyashtn   PetscCall(MatSetType(a->A, mtype));
13999566063dSJacob Faibussowitsch   PetscCall(MatSeqDenseSetPreallocation(a->A, data));
140047d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
14013faff063SStefano Zampini   mat->offloadmask = a->A->offloadmask;
14023faff063SStefano Zampini #endif
1403637a0070SStefano Zampini   mat->preallocated = PETSC_TRUE;
14046f08ad05SPierre Jolivet   mat->assembled    = PETSC_TRUE;
14053ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1406a23d5eceSKris Buschelman }
1407a23d5eceSKris Buschelman 
1408d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode MatConvert_MPIAIJ_MPIDense(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1409d71ae5a4SJacob Faibussowitsch {
14104e29119aSPierre Jolivet   Mat B, C;
14114e29119aSPierre Jolivet 
14124e29119aSPierre Jolivet   PetscFunctionBegin;
14139566063dSJacob Faibussowitsch   PetscCall(MatMPIAIJGetLocalMat(A, MAT_INITIAL_MATRIX, &C));
14149566063dSJacob Faibussowitsch   PetscCall(MatConvert_SeqAIJ_SeqDense(C, MATSEQDENSE, MAT_INITIAL_MATRIX, &B));
14159566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&C));
14164e29119aSPierre Jolivet   if (reuse == MAT_REUSE_MATRIX) {
14174e29119aSPierre Jolivet     C = *newmat;
14184e29119aSPierre Jolivet   } else C = NULL;
14199566063dSJacob Faibussowitsch   PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C));
14209566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&B));
1421ac530a7eSPierre Jolivet   if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &C));
1422ac530a7eSPierre Jolivet   else if (reuse == MAT_INITIAL_MATRIX) *newmat = C;
14233ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14244e29119aSPierre Jolivet }
14254e29119aSPierre Jolivet 
142666976f2fSJacob Faibussowitsch static PetscErrorCode MatConvert_MPIDense_MPIAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1427d71ae5a4SJacob Faibussowitsch {
14284e29119aSPierre Jolivet   Mat B, C;
14294e29119aSPierre Jolivet 
14304e29119aSPierre Jolivet   PetscFunctionBegin;
14319566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLocalMatrix(A, &C));
14329566063dSJacob Faibussowitsch   PetscCall(MatConvert_SeqDense_SeqAIJ(C, MATSEQAIJ, MAT_INITIAL_MATRIX, &B));
14334e29119aSPierre Jolivet   if (reuse == MAT_REUSE_MATRIX) {
14344e29119aSPierre Jolivet     C = *newmat;
14354e29119aSPierre Jolivet   } else C = NULL;
14369566063dSJacob Faibussowitsch   PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C));
14379566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&B));
1438ac530a7eSPierre Jolivet   if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &C));
1439ac530a7eSPierre Jolivet   else if (reuse == MAT_INITIAL_MATRIX) *newmat = C;
14403ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14414e29119aSPierre Jolivet }
14424e29119aSPierre Jolivet 
144365b80a83SHong Zhang #if defined(PETSC_HAVE_ELEMENTAL)
1444d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode MatConvert_MPIDense_Elemental(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1445d71ae5a4SJacob Faibussowitsch {
1446f709e8d7SPierre Jolivet   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
14478ea901baSHong Zhang   Mat           mat_elemental;
144832d7a744SHong Zhang   PetscScalar  *v;
1449f709e8d7SPierre Jolivet   PetscInt      m = A->rmap->n, N = A->cmap->N, rstart = A->rmap->rstart, i, *rows, *cols, lda;
14508ea901baSHong Zhang 
14518baccfbdSHong Zhang   PetscFunctionBegin;
1452378336b6SHong Zhang   if (reuse == MAT_REUSE_MATRIX) {
1453378336b6SHong Zhang     mat_elemental = *newmat;
14549566063dSJacob Faibussowitsch     PetscCall(MatZeroEntries(*newmat));
1455378336b6SHong Zhang   } else {
14569566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat_elemental));
14579566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(mat_elemental, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N));
14589566063dSJacob Faibussowitsch     PetscCall(MatSetType(mat_elemental, MATELEMENTAL));
14599566063dSJacob Faibussowitsch     PetscCall(MatSetUp(mat_elemental));
14609566063dSJacob Faibussowitsch     PetscCall(MatSetOption(mat_elemental, MAT_ROW_ORIENTED, PETSC_FALSE));
1461378336b6SHong Zhang   }
1462378336b6SHong Zhang 
14639566063dSJacob Faibussowitsch   PetscCall(PetscMalloc2(m, &rows, N, &cols));
146432d7a744SHong Zhang   for (i = 0; i < N; i++) cols[i] = i;
146532d7a744SHong Zhang   for (i = 0; i < m; i++) rows[i] = rstart + i;
14668ea901baSHong Zhang 
1467637a0070SStefano Zampini   /* PETSc-Elemental interface uses axpy for setting off-processor entries, only ADD_VALUES is allowed */
14689566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(A, &v));
1469f709e8d7SPierre Jolivet   PetscCall(MatDenseGetLDA(a->A, &lda));
1470f709e8d7SPierre Jolivet   if (lda == m) PetscCall(MatSetValues(mat_elemental, m, rows, N, cols, v, ADD_VALUES));
1471f709e8d7SPierre Jolivet   else {
1472f709e8d7SPierre Jolivet     for (i = 0; i < N; i++) PetscCall(MatSetValues(mat_elemental, m, rows, 1, &i, v + lda * i, ADD_VALUES));
1473f709e8d7SPierre Jolivet   }
14749566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(mat_elemental, MAT_FINAL_ASSEMBLY));
14759566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(mat_elemental, MAT_FINAL_ASSEMBLY));
14769566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(A, &v));
14779566063dSJacob Faibussowitsch   PetscCall(PetscFree2(rows, cols));
14788ea901baSHong Zhang 
1479511c6705SHong Zhang   if (reuse == MAT_INPLACE_MATRIX) {
14809566063dSJacob Faibussowitsch     PetscCall(MatHeaderReplace(A, &mat_elemental));
14818ea901baSHong Zhang   } else {
14828ea901baSHong Zhang     *newmat = mat_elemental;
14838ea901baSHong Zhang   }
14843ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14858baccfbdSHong Zhang }
148665b80a83SHong Zhang #endif
14878baccfbdSHong Zhang 
1488d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetColumn_MPIDense(Mat A, PetscInt col, PetscScalar **vals)
1489d71ae5a4SJacob Faibussowitsch {
149086aefd0dSHong Zhang   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
149186aefd0dSHong Zhang 
149286aefd0dSHong Zhang   PetscFunctionBegin;
14939566063dSJacob Faibussowitsch   PetscCall(MatDenseGetColumn(mat->A, col, vals));
14943ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
149586aefd0dSHong Zhang }
149686aefd0dSHong Zhang 
1497d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseRestoreColumn_MPIDense(Mat A, PetscScalar **vals)
1498d71ae5a4SJacob Faibussowitsch {
149986aefd0dSHong Zhang   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
150086aefd0dSHong Zhang 
150186aefd0dSHong Zhang   PetscFunctionBegin;
15029566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreColumn(mat->A, vals));
15033ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
150486aefd0dSHong Zhang }
150586aefd0dSHong Zhang 
1506d71ae5a4SJacob Faibussowitsch PetscErrorCode MatCreateMPIMatConcatenateSeqMat_MPIDense(MPI_Comm comm, Mat inmat, PetscInt n, MatReuse scall, Mat *outmat)
1507d71ae5a4SJacob Faibussowitsch {
150894e2cb23SJakub Kruzik   Mat_MPIDense *mat;
150994e2cb23SJakub Kruzik   PetscInt      m, nloc, N;
151094e2cb23SJakub Kruzik 
151194e2cb23SJakub Kruzik   PetscFunctionBegin;
15129566063dSJacob Faibussowitsch   PetscCall(MatGetSize(inmat, &m, &N));
15139566063dSJacob Faibussowitsch   PetscCall(MatGetLocalSize(inmat, NULL, &nloc));
151494e2cb23SJakub Kruzik   if (scall == MAT_INITIAL_MATRIX) { /* symbolic phase */
151594e2cb23SJakub Kruzik     PetscInt sum;
151694e2cb23SJakub Kruzik 
151748a46eb9SPierre Jolivet     if (n == PETSC_DECIDE) PetscCall(PetscSplitOwnership(comm, &n, &N));
151894e2cb23SJakub Kruzik     /* Check sum(n) = N */
1519462c564dSBarry Smith     PetscCallMPI(MPIU_Allreduce(&n, &sum, 1, MPIU_INT, MPI_SUM, comm));
152008401ef6SPierre Jolivet     PetscCheck(sum == N, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Sum of local columns %" PetscInt_FMT " != global columns %" PetscInt_FMT, sum, N);
152194e2cb23SJakub Kruzik 
15229566063dSJacob Faibussowitsch     PetscCall(MatCreateDense(comm, m, n, PETSC_DETERMINE, N, NULL, outmat));
15239566063dSJacob Faibussowitsch     PetscCall(MatSetOption(*outmat, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
152494e2cb23SJakub Kruzik   }
152594e2cb23SJakub Kruzik 
152694e2cb23SJakub Kruzik   /* numeric phase */
152794e2cb23SJakub Kruzik   mat = (Mat_MPIDense *)(*outmat)->data;
15289566063dSJacob Faibussowitsch   PetscCall(MatCopy(inmat, mat->A, SAME_NONZERO_PATTERN));
15293ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
153094e2cb23SJakub Kruzik }
153194e2cb23SJakub Kruzik 
1532d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVec_MPIDense(Mat A, PetscInt col, Vec *v)
1533d71ae5a4SJacob Faibussowitsch {
15346947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15356947451fSStefano Zampini   PetscInt      lda;
15366947451fSStefano Zampini 
15376947451fSStefano Zampini   PetscFunctionBegin;
153828b400f6SJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
153928b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1540d16ceb75SStefano Zampini   if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
15416947451fSStefano Zampini   a->vecinuse = col + 1;
15429566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
15439566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(a->A, (PetscScalar **)&a->ptrinuse));
15448e3a54c0SPierre Jolivet   PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
15456947451fSStefano Zampini   *v = a->cvec;
15463ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15476947451fSStefano Zampini }
15486947451fSStefano Zampini 
1549d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVec_MPIDense(Mat A, PetscInt col, Vec *v)
1550d71ae5a4SJacob Faibussowitsch {
15516947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15526947451fSStefano Zampini 
15536947451fSStefano Zampini   PetscFunctionBegin;
155428b400f6SJacob Faibussowitsch   PetscCheck(a->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
155528b400f6SJacob Faibussowitsch   PetscCheck(a->cvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column vector");
15564186a4bbSPierre Jolivet   VecCheckAssembled(a->cvec);
15576947451fSStefano Zampini   a->vecinuse = 0;
15589566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(a->A, (PetscScalar **)&a->ptrinuse));
15599566063dSJacob Faibussowitsch   PetscCall(VecResetArray(a->cvec));
1560742765d3SMatthew Knepley   if (v) *v = NULL;
15613ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15626947451fSStefano Zampini }
15636947451fSStefano Zampini 
1564d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v)
1565d71ae5a4SJacob Faibussowitsch {
15666947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15676947451fSStefano Zampini   PetscInt      lda;
15686947451fSStefano Zampini 
15696947451fSStefano Zampini   PetscFunctionBegin;
157028b400f6SJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
157128b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1572d16ceb75SStefano Zampini   if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
15736947451fSStefano Zampini   a->vecinuse = col + 1;
15749566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
15759566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, &a->ptrinuse));
15768e3a54c0SPierre Jolivet   PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
15779566063dSJacob Faibussowitsch   PetscCall(VecLockReadPush(a->cvec));
15786947451fSStefano Zampini   *v = a->cvec;
15793ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15806947451fSStefano Zampini }
15816947451fSStefano Zampini 
1582d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v)
1583d71ae5a4SJacob Faibussowitsch {
15846947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15856947451fSStefano Zampini 
15866947451fSStefano Zampini   PetscFunctionBegin;
15875f80ce2aSJacob Faibussowitsch   PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
15885f80ce2aSJacob Faibussowitsch   PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector");
15894186a4bbSPierre Jolivet   VecCheckAssembled(a->cvec);
15906947451fSStefano Zampini   a->vecinuse = 0;
15919566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, &a->ptrinuse));
15929566063dSJacob Faibussowitsch   PetscCall(VecLockReadPop(a->cvec));
15939566063dSJacob Faibussowitsch   PetscCall(VecResetArray(a->cvec));
1594742765d3SMatthew Knepley   if (v) *v = NULL;
15953ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15966947451fSStefano Zampini }
15976947451fSStefano Zampini 
1598d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v)
1599d71ae5a4SJacob Faibussowitsch {
16006947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
16016947451fSStefano Zampini   PetscInt      lda;
16026947451fSStefano Zampini 
16036947451fSStefano Zampini   PetscFunctionBegin;
16045f80ce2aSJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
16055f80ce2aSJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1606d16ceb75SStefano Zampini   if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
16076947451fSStefano Zampini   a->vecinuse = col + 1;
16089566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
16099566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(a->A, (PetscScalar **)&a->ptrinuse));
16108e3a54c0SPierre Jolivet   PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
16116947451fSStefano Zampini   *v = a->cvec;
16123ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
16136947451fSStefano Zampini }
16146947451fSStefano Zampini 
1615d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v)
1616d71ae5a4SJacob Faibussowitsch {
16176947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
16186947451fSStefano Zampini 
16196947451fSStefano Zampini   PetscFunctionBegin;
16205f80ce2aSJacob Faibussowitsch   PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
16215f80ce2aSJacob Faibussowitsch   PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector");
16224186a4bbSPierre Jolivet   VecCheckAssembled(a->cvec);
16236947451fSStefano Zampini   a->vecinuse = 0;
16249566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(a->A, (PetscScalar **)&a->ptrinuse));
16259566063dSJacob Faibussowitsch   PetscCall(VecResetArray(a->cvec));
1626742765d3SMatthew Knepley   if (v) *v = NULL;
16273ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
16286947451fSStefano Zampini }
16296947451fSStefano Zampini 
163066976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseGetSubMatrix_MPIDense(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *v)
1631d71ae5a4SJacob Faibussowitsch {
16325ea7661aSPierre Jolivet   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1633616b8fbbSStefano Zampini   Mat_MPIDense *c;
1634616b8fbbSStefano Zampini   MPI_Comm      comm;
1635d016bddeSToby Isaac   PetscInt      prbegin, prend, pcbegin, pcend;
16365ea7661aSPierre Jolivet 
16375ea7661aSPierre Jolivet   PetscFunctionBegin;
16389566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
16395f80ce2aSJacob Faibussowitsch   PetscCheck(!a->vecinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
16405f80ce2aSJacob Faibussowitsch   PetscCheck(!a->matinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1641d016bddeSToby Isaac   prbegin = PetscMax(0, PetscMin(A->rmap->rend, rbegin) - A->rmap->rstart);
1642d016bddeSToby Isaac   prend   = PetscMin(A->rmap->n, PetscMax(0, rend - A->rmap->rstart));
1643d016bddeSToby Isaac   pcbegin = PetscMax(0, PetscMin(A->cmap->rend, cbegin) - A->cmap->rstart);
1644d016bddeSToby Isaac   pcend   = PetscMin(A->cmap->n, PetscMax(0, cend - A->cmap->rstart));
16455ea7661aSPierre Jolivet   if (!a->cmat) {
16469566063dSJacob Faibussowitsch     PetscCall(MatCreate(comm, &a->cmat));
16479566063dSJacob Faibussowitsch     PetscCall(MatSetType(a->cmat, ((PetscObject)A)->type_name));
1648a2748737SPierre Jolivet     if (rend - rbegin == A->rmap->N) PetscCall(PetscLayoutReference(A->rmap, &a->cmat->rmap));
1649a2748737SPierre Jolivet     else {
1650d016bddeSToby Isaac       PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, prend - prbegin));
1651a2748737SPierre Jolivet       PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin));
1652a2748737SPierre Jolivet       PetscCall(PetscLayoutSetUp(a->cmat->rmap));
1653a2748737SPierre Jolivet     }
1654d016bddeSToby Isaac     if (cend - cbegin == A->cmap->N) PetscCall(PetscLayoutReference(A->cmap, &a->cmat->cmap));
1655d016bddeSToby Isaac     else {
1656d016bddeSToby Isaac       PetscCall(PetscLayoutSetLocalSize(a->cmat->cmap, pcend - pcbegin));
16579566063dSJacob Faibussowitsch       PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin));
16589566063dSJacob Faibussowitsch       PetscCall(PetscLayoutSetUp(a->cmat->cmap));
1659a2748737SPierre Jolivet     }
1660616b8fbbSStefano Zampini     c             = (Mat_MPIDense *)a->cmat->data;
1661d016bddeSToby Isaac     c->sub_rbegin = rbegin;
1662d016bddeSToby Isaac     c->sub_rend   = rend;
1663d016bddeSToby Isaac     c->sub_cbegin = cbegin;
1664d016bddeSToby Isaac     c->sub_cend   = cend;
1665d016bddeSToby Isaac   }
1666d016bddeSToby Isaac   c = (Mat_MPIDense *)a->cmat->data;
1667d016bddeSToby Isaac   if (c->sub_rbegin != rbegin || c->sub_rend != rend) {
1668d016bddeSToby Isaac     PetscCall(PetscLayoutDestroy(&a->cmat->rmap));
1669d016bddeSToby Isaac     PetscCall(PetscLayoutCreate(comm, &a->cmat->rmap));
1670d016bddeSToby Isaac     PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, prend - prbegin));
1671d016bddeSToby Isaac     PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin));
1672d016bddeSToby Isaac     PetscCall(PetscLayoutSetUp(a->cmat->rmap));
1673d016bddeSToby Isaac     c->sub_rbegin = rbegin;
1674d016bddeSToby Isaac     c->sub_rend   = rend;
1675d016bddeSToby Isaac   }
1676d016bddeSToby Isaac   if (c->sub_cbegin != cbegin || c->sub_cend != cend) {
1677d016bddeSToby Isaac     // special optimization: check if all columns are owned by rank 0, in which case no communication is necessary
1678d016bddeSToby Isaac     if ((cend - cbegin != a->cmat->cmap->N) || (A->cmap->range[1] != A->cmap->N)) {
1679d016bddeSToby Isaac       PetscCall(PetscLayoutDestroy(&a->cmat->cmap));
1680d016bddeSToby Isaac       PetscCall(PetscLayoutCreate(comm, &a->cmat->cmap));
1681d016bddeSToby Isaac       PetscCall(PetscLayoutSetLocalSize(a->cmat->cmap, pcend - pcbegin));
1682d016bddeSToby Isaac       PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin));
1683d016bddeSToby Isaac       PetscCall(PetscLayoutSetUp(a->cmat->cmap));
1684d016bddeSToby Isaac       PetscCall(VecDestroy(&c->lvec));
1685d016bddeSToby Isaac       PetscCall(PetscSFDestroy(&c->Mvctx));
1686d016bddeSToby Isaac     }
1687d016bddeSToby Isaac     c->sub_cbegin = cbegin;
1688d016bddeSToby Isaac     c->sub_cend   = cend;
1689d016bddeSToby Isaac   }
16905f80ce2aSJacob Faibussowitsch   PetscCheck(!c->A, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1691d016bddeSToby Isaac   PetscCall(MatDenseGetSubMatrix(a->A, prbegin, prend, cbegin, cend, &c->A));
16923faff063SStefano Zampini 
1693616b8fbbSStefano Zampini   a->cmat->preallocated = PETSC_TRUE;
1694616b8fbbSStefano Zampini   a->cmat->assembled    = PETSC_TRUE;
16953faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE)
16963faff063SStefano Zampini   a->cmat->offloadmask = c->A->offloadmask;
16973faff063SStefano Zampini #endif
16985ea7661aSPierre Jolivet   a->matinuse = cbegin + 1;
16995ea7661aSPierre Jolivet   *v          = a->cmat;
17003ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
17015ea7661aSPierre Jolivet }
17025ea7661aSPierre Jolivet 
170366976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreSubMatrix_MPIDense(Mat A, Mat *v)
1704d71ae5a4SJacob Faibussowitsch {
17055ea7661aSPierre Jolivet   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1706616b8fbbSStefano Zampini   Mat_MPIDense *c;
17075ea7661aSPierre Jolivet 
17085ea7661aSPierre Jolivet   PetscFunctionBegin;
17095f80ce2aSJacob Faibussowitsch   PetscCheck(a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetSubMatrix() first");
17105f80ce2aSJacob Faibussowitsch   PetscCheck(a->cmat, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal matrix");
17115f80ce2aSJacob Faibussowitsch   PetscCheck(*v == a->cmat, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Not the matrix obtained from MatDenseGetSubMatrix()");
17125ea7661aSPierre Jolivet   a->matinuse = 0;
1713616b8fbbSStefano Zampini   c           = (Mat_MPIDense *)a->cmat->data;
17149566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreSubMatrix(a->A, &c->A));
1715742765d3SMatthew Knepley   if (v) *v = NULL;
17163faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE)
17173faff063SStefano Zampini   A->offloadmask = a->A->offloadmask;
17183faff063SStefano Zampini #endif
17193ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
17205ea7661aSPierre Jolivet }
17215ea7661aSPierre Jolivet 
1722002bc6daSRichard Tran Mills /*MC
1723002bc6daSRichard Tran Mills    MATMPIDENSE - MATMPIDENSE = "mpidense" - A matrix type to be used for distributed dense matrices.
1724002bc6daSRichard Tran Mills 
17252ef1f0ffSBarry Smith    Options Database Key:
172611a5261eSBarry Smith . -mat_type mpidense - sets the matrix type to `MATMPIDENSE` during a call to `MatSetFromOptions()`
1727002bc6daSRichard Tran Mills 
1728002bc6daSRichard Tran Mills   Level: beginner
1729002bc6daSRichard Tran Mills 
17301cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MatCreateDense()`, `MATSEQDENSE`, `MATDENSE`
1731002bc6daSRichard Tran Mills M*/
17324742e46bSJacob Faibussowitsch PetscErrorCode MatCreate_MPIDense(Mat mat)
1733d71ae5a4SJacob Faibussowitsch {
1734273d9f13SBarry Smith   Mat_MPIDense *a;
1735273d9f13SBarry Smith 
1736273d9f13SBarry Smith   PetscFunctionBegin;
17374dfa11a4SJacob Faibussowitsch   PetscCall(PetscNew(&a));
1738b0a32e0cSBarry Smith   mat->data   = (void *)a;
1739aea10558SJacob Faibussowitsch   mat->ops[0] = MatOps_Values;
1740273d9f13SBarry Smith 
1741273d9f13SBarry Smith   mat->insertmode = NOT_SET_VALUES;
1742273d9f13SBarry Smith 
1743273d9f13SBarry Smith   /* build cache for off array entries formed */
1744273d9f13SBarry Smith   a->donotstash = PETSC_FALSE;
17452205254eSKarl Rupp 
17469566063dSJacob Faibussowitsch   PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)mat), 1, &mat->stash));
1747273d9f13SBarry Smith 
1748273d9f13SBarry Smith   /* stuff used for matrix vector multiply */
1749f4259b30SLisandro Dalcin   a->lvec        = NULL;
1750f4259b30SLisandro Dalcin   a->Mvctx       = NULL;
1751273d9f13SBarry Smith   a->roworiented = PETSC_TRUE;
1752273d9f13SBarry Smith 
17539566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", MatDenseGetLDA_MPIDense));
17549566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", MatDenseSetLDA_MPIDense));
17559566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", MatDenseGetArray_MPIDense));
17569566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", MatDenseRestoreArray_MPIDense));
17579566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", MatDenseGetArrayRead_MPIDense));
17589566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", MatDenseRestoreArrayRead_MPIDense));
17599566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", MatDenseGetArrayWrite_MPIDense));
17609566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", MatDenseRestoreArrayWrite_MPIDense));
17619566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", MatDensePlaceArray_MPIDense));
17629566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", MatDenseResetArray_MPIDense));
17639566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", MatDenseReplaceArray_MPIDense));
17649566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_MPIDense));
17659566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_MPIDense));
17669566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_MPIDense));
17679566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_MPIDense));
17689566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_MPIDense));
17699566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_MPIDense));
17709566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_MPIDense));
17719566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_MPIDense));
17729566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", MatConvert_MPIAIJ_MPIDense));
17739566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", MatConvert_MPIDense_MPIAIJ));
17748baccfbdSHong Zhang #if defined(PETSC_HAVE_ELEMENTAL)
17759566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", MatConvert_MPIDense_Elemental));
17768baccfbdSHong Zhang #endif
1777*d1a032dbSPierre Jolivet #if defined(PETSC_HAVE_SCALAPACK) && (defined(PETSC_USE_REAL_SINGLE) || defined(PETSC_USE_REAL_DOUBLE))
17789566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", MatConvert_Dense_ScaLAPACK));
1779d24d4204SJose E. Roman #endif
1780637a0070SStefano Zampini #if defined(PETSC_HAVE_CUDA)
17819566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", MatConvert_MPIDense_MPIDenseCUDA));
1782637a0070SStefano Zampini #endif
17839566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", MatMPIDenseSetPreallocation_MPIDense));
17849566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
17859566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
17866718818eSStefano Zampini #if defined(PETSC_HAVE_CUDA)
17879566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
17889566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
17896718818eSStefano Zampini #endif
179047d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
179147d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", MatConvert_MPIDense_MPIDenseHIP));
179247d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
179347d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
179447d993e7Ssuyashtn #endif
17959566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", MatDenseGetColumn_MPIDense));
17969566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", MatDenseRestoreColumn_MPIDense));
1797d016bddeSToby Isaac   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultColumnRange_C", MatMultColumnRange_MPIDense));
17980be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", MatMultAddColumnRange_MPIDense));
17990be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", MatMultHermitianTransposeColumnRange_MPIDense));
18000be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", MatMultHermitianTransposeAddColumnRange_MPIDense));
18019566063dSJacob Faibussowitsch   PetscCall(PetscObjectChangeTypeName((PetscObject)mat, MATMPIDENSE));
18023ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1803273d9f13SBarry Smith }
1804273d9f13SBarry Smith 
1805209238afSKris Buschelman /*MC
1806002d173eSKris Buschelman    MATDENSE - MATDENSE = "dense" - A matrix type to be used for dense matrices.
1807209238afSKris Buschelman 
180811a5261eSBarry Smith    This matrix type is identical to `MATSEQDENSE` when constructed with a single process communicator,
180911a5261eSBarry Smith    and `MATMPIDENSE` otherwise.
1810209238afSKris Buschelman 
18112ef1f0ffSBarry Smith    Options Database Key:
181211a5261eSBarry Smith . -mat_type dense - sets the matrix type to `MATDENSE` during a call to `MatSetFromOptions()`
1813209238afSKris Buschelman 
1814209238afSKris Buschelman   Level: beginner
1815209238afSKris Buschelman 
18161cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATSEQDENSE`, `MATMPIDENSE`, `MATDENSECUDA`, `MATDENSEHIP`
18176947451fSStefano Zampini M*/
18186947451fSStefano Zampini 
18195d83a8b1SBarry Smith /*@
1820273d9f13SBarry Smith   MatMPIDenseSetPreallocation - Sets the array used to store the matrix entries
1821273d9f13SBarry Smith 
18228f490ea3SStefano Zampini   Collective
1823273d9f13SBarry Smith 
1824273d9f13SBarry Smith   Input Parameters:
1825fe59aa6dSJacob Faibussowitsch + B    - the matrix
18262ef1f0ffSBarry Smith - data - optional location of matrix data.  Set to `NULL` for PETSc
1827273d9f13SBarry Smith          to control all matrix memory allocation.
1828273d9f13SBarry Smith 
18292ef1f0ffSBarry Smith   Level: intermediate
18302ef1f0ffSBarry Smith 
1831273d9f13SBarry Smith   Notes:
18322ef1f0ffSBarry Smith   The dense format is fully compatible with standard Fortran
1833273d9f13SBarry Smith   storage by columns.
1834273d9f13SBarry Smith 
1835273d9f13SBarry Smith   The data input variable is intended primarily for Fortran programmers
1836273d9f13SBarry Smith   who wish to allocate their own matrix memory space.  Most users should
18372ef1f0ffSBarry Smith   set `data` to `NULL`.
1838273d9f13SBarry Smith 
18391cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATMPIDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()`
1840273d9f13SBarry Smith @*/
1841d71ae5a4SJacob Faibussowitsch PetscErrorCode MatMPIDenseSetPreallocation(Mat B, PetscScalar *data)
1842d71ae5a4SJacob Faibussowitsch {
1843273d9f13SBarry Smith   PetscFunctionBegin;
1844d5ea218eSStefano Zampini   PetscValidHeaderSpecific(B, MAT_CLASSID, 1);
1845cac4c232SBarry Smith   PetscTryMethod(B, "MatMPIDenseSetPreallocation_C", (Mat, PetscScalar *), (B, data));
18463ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1847273d9f13SBarry Smith }
1848273d9f13SBarry Smith 
1849d3042a70SBarry Smith /*@
185011a5261eSBarry Smith   MatDensePlaceArray - Allows one to replace the array in a `MATDENSE` matrix with an
1851d3042a70SBarry Smith   array provided by the user. This is useful to avoid copying an array
1852d3042a70SBarry Smith   into a matrix
1853d3042a70SBarry Smith 
1854d3042a70SBarry Smith   Not Collective
1855d3042a70SBarry Smith 
1856d3042a70SBarry Smith   Input Parameters:
1857d3042a70SBarry Smith + mat   - the matrix
1858d3042a70SBarry Smith - array - the array in column major order
1859d3042a70SBarry Smith 
18602ef1f0ffSBarry Smith   Level: developer
18612ef1f0ffSBarry Smith 
186211a5261eSBarry Smith   Note:
18634d5500e8SJunchao Zhang   Adding `const` to `array` was an oversight, see notes in `VecPlaceArray()`.
18644d5500e8SJunchao Zhang 
186511a5261eSBarry 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
1866d3042a70SBarry Smith   freed when the matrix is destroyed.
1867d3042a70SBarry Smith 
18681cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDenseResetArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`,
186911a5261eSBarry Smith           `MatDenseReplaceArray()`
1870d3042a70SBarry Smith @*/
1871d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDensePlaceArray(Mat mat, const PetscScalar *array)
1872d71ae5a4SJacob Faibussowitsch {
1873d3042a70SBarry Smith   PetscFunctionBegin;
1874d5ea218eSStefano Zampini   PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
1875cac4c232SBarry Smith   PetscUseMethod(mat, "MatDensePlaceArray_C", (Mat, const PetscScalar *), (mat, array));
18769566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)mat));
187747d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
1878637a0070SStefano Zampini   mat->offloadmask = PETSC_OFFLOAD_CPU;
1879637a0070SStefano Zampini #endif
18803ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1881d3042a70SBarry Smith }
1882d3042a70SBarry Smith 
1883d3042a70SBarry Smith /*@
188411a5261eSBarry Smith   MatDenseResetArray - Resets the matrix array to that it previously had before the call to `MatDensePlaceArray()`
1885d3042a70SBarry Smith 
1886d3042a70SBarry Smith   Not Collective
1887d3042a70SBarry Smith 
18882fe279fdSBarry Smith   Input Parameter:
1889d3042a70SBarry Smith . mat - the matrix
1890d3042a70SBarry Smith 
18912ef1f0ffSBarry Smith   Level: developer
18922ef1f0ffSBarry Smith 
189311a5261eSBarry Smith   Note:
189411a5261eSBarry Smith   You can only call this after a call to `MatDensePlaceArray()`
1895d3042a70SBarry Smith 
18961cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDensePlaceArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`
1897d3042a70SBarry Smith @*/
1898d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseResetArray(Mat mat)
1899d71ae5a4SJacob Faibussowitsch {
1900d3042a70SBarry Smith   PetscFunctionBegin;
1901d5ea218eSStefano Zampini   PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
1902cac4c232SBarry Smith   PetscUseMethod(mat, "MatDenseResetArray_C", (Mat), (mat));
19039566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)mat));
19043ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1905d3042a70SBarry Smith }
1906d3042a70SBarry Smith 
1907d5ea218eSStefano Zampini /*@
1908d5ea218eSStefano Zampini   MatDenseReplaceArray - Allows one to replace the array in a dense matrix with an
1909d5ea218eSStefano Zampini   array provided by the user. This is useful to avoid copying an array
1910d5ea218eSStefano Zampini   into a matrix
1911d5ea218eSStefano Zampini 
1912d5ea218eSStefano Zampini   Not Collective
1913d5ea218eSStefano Zampini 
1914d5ea218eSStefano Zampini   Input Parameters:
1915d5ea218eSStefano Zampini + mat   - the matrix
1916d5ea218eSStefano Zampini - array - the array in column major order
1917d5ea218eSStefano Zampini 
19182ef1f0ffSBarry Smith   Level: developer
19192ef1f0ffSBarry Smith 
192011a5261eSBarry Smith   Note:
19214d5500e8SJunchao Zhang   Adding `const` to `array` was an oversight, see notes in `VecPlaceArray()`.
19224d5500e8SJunchao Zhang 
192311a5261eSBarry Smith   The memory passed in MUST be obtained with `PetscMalloc()` and CANNOT be
1924d5ea218eSStefano Zampini   freed by the user. It will be freed when the matrix is destroyed.
1925d5ea218eSStefano Zampini 
19261cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MatDensePlaceArray()`, `MatDenseGetArray()`, `VecReplaceArray()`
1927d5ea218eSStefano Zampini @*/
1928d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseReplaceArray(Mat mat, const PetscScalar *array)
1929d71ae5a4SJacob Faibussowitsch {
1930d5ea218eSStefano Zampini   PetscFunctionBegin;
1931d5ea218eSStefano Zampini   PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
1932cac4c232SBarry Smith   PetscUseMethod(mat, "MatDenseReplaceArray_C", (Mat, const PetscScalar *), (mat, array));
19339566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)mat));
193447d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
1935d5ea218eSStefano Zampini   mat->offloadmask = PETSC_OFFLOAD_CPU;
1936d5ea218eSStefano Zampini #endif
19373ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1938d5ea218eSStefano Zampini }
1939d5ea218eSStefano Zampini 
19405d83a8b1SBarry Smith /*@
194111a5261eSBarry Smith   MatCreateDense - Creates a matrix in `MATDENSE` format.
19428965ea79SLois Curfman McInnes 
1943d083f849SBarry Smith   Collective
1944db81eaa0SLois Curfman McInnes 
19458965ea79SLois Curfman McInnes   Input Parameters:
1946db81eaa0SLois Curfman McInnes + comm - MPI communicator
19472ef1f0ffSBarry Smith . m    - number of local rows (or `PETSC_DECIDE` to have calculated if `M` is given)
19482ef1f0ffSBarry Smith . n    - number of local columns (or `PETSC_DECIDE` to have calculated if `N` is given)
19492ef1f0ffSBarry Smith . M    - number of global rows (or `PETSC_DECIDE` to have calculated if `m` is given)
19502ef1f0ffSBarry Smith . N    - number of global columns (or `PETSC_DECIDE` to have calculated if `n` is given)
19512c5cda21SPierre Jolivet - data - optional location of matrix data.  Set data to `NULL` (`PETSC_NULL_SCALAR_ARRAY` for Fortran users) for PETSc
1952dfc5480cSLois Curfman McInnes    to control all matrix memory allocation.
19538965ea79SLois Curfman McInnes 
19548965ea79SLois Curfman McInnes   Output Parameter:
1955477f1c0bSLois Curfman McInnes . A - the matrix
19568965ea79SLois Curfman McInnes 
19572ef1f0ffSBarry Smith   Level: intermediate
19582ef1f0ffSBarry Smith 
1959b259b22eSLois Curfman McInnes   Notes:
19602ef1f0ffSBarry Smith   The dense format is fully compatible with standard Fortran
196139ddd567SLois Curfman McInnes   storage by columns.
196211a5261eSBarry Smith 
196311a5261eSBarry Smith   Although local portions of the matrix are stored in column-major
1964002bc6daSRichard Tran Mills   order, the matrix is partitioned across MPI ranks by row.
19658965ea79SLois Curfman McInnes 
196618f449edSLois Curfman McInnes   The data input variable is intended primarily for Fortran programmers
196718f449edSLois Curfman McInnes   who wish to allocate their own matrix memory space.  Most users should
19682c5cda21SPierre Jolivet   set `data` to `NULL` (`PETSC_NULL_SCALAR_ARRAY` for Fortran users).
196918f449edSLois Curfman McInnes 
19708965ea79SLois Curfman McInnes   The user MUST specify either the local or global matrix dimensions
19718965ea79SLois Curfman McInnes   (possibly both).
19728965ea79SLois Curfman McInnes 
19731cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()`
19748965ea79SLois Curfman McInnes @*/
1975ce78bad3SBarry Smith PetscErrorCode MatCreateDense(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt M, PetscInt N, PetscScalar data[], Mat *A)
1976d71ae5a4SJacob Faibussowitsch {
19773a40ed3dSBarry Smith   PetscFunctionBegin;
19789566063dSJacob Faibussowitsch   PetscCall(MatCreate(comm, A));
19799566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(*A, m, n, M, N));
19806f08ad05SPierre Jolivet   PetscCall(MatSetType(*A, MATDENSE));
19819566063dSJacob Faibussowitsch   PetscCall(MatSeqDenseSetPreallocation(*A, data));
19826f08ad05SPierre Jolivet   PetscCall(MatMPIDenseSetPreallocation(*A, data));
19833ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
19848965ea79SLois Curfman McInnes }
19858965ea79SLois Curfman McInnes 
1986d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDuplicate_MPIDense(Mat A, MatDuplicateOption cpvalues, Mat *newmat)
1987d71ae5a4SJacob Faibussowitsch {
19888965ea79SLois Curfman McInnes   Mat           mat;
19893501a2bdSLois Curfman McInnes   Mat_MPIDense *a, *oldmat = (Mat_MPIDense *)A->data;
19908965ea79SLois Curfman McInnes 
19913a40ed3dSBarry Smith   PetscFunctionBegin;
1992f4259b30SLisandro Dalcin   *newmat = NULL;
19939566063dSJacob Faibussowitsch   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat));
19949566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(mat, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N));
19959566063dSJacob Faibussowitsch   PetscCall(MatSetType(mat, ((PetscObject)A)->type_name));
1996834f8fabSBarry Smith   a = (Mat_MPIDense *)mat->data;
19975aa7edbeSHong Zhang 
1998d5f3da31SBarry Smith   mat->factortype   = A->factortype;
1999c456f294SBarry Smith   mat->assembled    = PETSC_TRUE;
2000273d9f13SBarry Smith   mat->preallocated = PETSC_TRUE;
20018965ea79SLois Curfman McInnes 
2002e0fa3b82SLois Curfman McInnes   mat->insertmode = NOT_SET_VALUES;
20033782ba37SSatish Balay   a->donotstash   = oldmat->donotstash;
2004e04c1aa4SHong Zhang 
20059566063dSJacob Faibussowitsch   PetscCall(PetscLayoutReference(A->rmap, &mat->rmap));
20069566063dSJacob Faibussowitsch   PetscCall(PetscLayoutReference(A->cmap, &mat->cmap));
20078965ea79SLois Curfman McInnes 
20089566063dSJacob Faibussowitsch   PetscCall(MatDuplicate(oldmat->A, cpvalues, &a->A));
200901b82886SBarry Smith 
20108965ea79SLois Curfman McInnes   *newmat = mat;
20113ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
20128965ea79SLois Curfman McInnes }
20138965ea79SLois Curfman McInnes 
201466976f2fSJacob Faibussowitsch static PetscErrorCode MatLoad_MPIDense(Mat newMat, PetscViewer viewer)
2015d71ae5a4SJacob Faibussowitsch {
201687d5ce66SSatish Balay   PetscBool isbinary;
201787d5ce66SSatish Balay #if defined(PETSC_HAVE_HDF5)
201887d5ce66SSatish Balay   PetscBool ishdf5;
201987d5ce66SSatish Balay #endif
2020eb91f321SVaclav Hapla 
2021eb91f321SVaclav Hapla   PetscFunctionBegin;
2022eb91f321SVaclav Hapla   PetscValidHeaderSpecific(newMat, MAT_CLASSID, 1);
2023eb91f321SVaclav Hapla   PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
2024eb91f321SVaclav Hapla   /* force binary viewer to load .info file if it has not yet done so */
20259566063dSJacob Faibussowitsch   PetscCall(PetscViewerSetUp(viewer));
20269566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
202787d5ce66SSatish Balay #if defined(PETSC_HAVE_HDF5)
20289566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERHDF5, &ishdf5));
202987d5ce66SSatish Balay #endif
2030eb91f321SVaclav Hapla   if (isbinary) {
20319566063dSJacob Faibussowitsch     PetscCall(MatLoad_Dense_Binary(newMat, viewer));
2032eb91f321SVaclav Hapla #if defined(PETSC_HAVE_HDF5)
203387d5ce66SSatish Balay   } else if (ishdf5) {
20349566063dSJacob Faibussowitsch     PetscCall(MatLoad_Dense_HDF5(newMat, viewer));
2035eb91f321SVaclav Hapla #endif
203698921bdaSJacob 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);
20373ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2038eb91f321SVaclav Hapla }
2039eb91f321SVaclav Hapla 
2040d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatEqual_MPIDense(Mat A, Mat B, PetscBool *flag)
2041d71ae5a4SJacob Faibussowitsch {
20426e4ee0c6SHong Zhang   Mat_MPIDense *matB = (Mat_MPIDense *)B->data, *matA = (Mat_MPIDense *)A->data;
20436e4ee0c6SHong Zhang   Mat           a, b;
204490ace30eSBarry Smith 
20456e4ee0c6SHong Zhang   PetscFunctionBegin;
20466e4ee0c6SHong Zhang   a = matA->A;
20476e4ee0c6SHong Zhang   b = matB->A;
20482cf15c64SPierre Jolivet   PetscCall(MatEqual(a, b, flag));
20495440e5dcSBarry Smith   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, flag, 1, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
20503ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
20516e4ee0c6SHong Zhang }
205290ace30eSBarry Smith 
205366976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MatTransMatMult_MPIDense_MPIDense(void *data)
2054d71ae5a4SJacob Faibussowitsch {
20556718818eSStefano Zampini   Mat_TransMatMultDense *atb = (Mat_TransMatMultDense *)data;
2056baa3c1c6SHong Zhang 
2057baa3c1c6SHong Zhang   PetscFunctionBegin;
20589566063dSJacob Faibussowitsch   PetscCall(PetscFree2(atb->sendbuf, atb->recvcounts));
20599566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&atb->atb));
20609566063dSJacob Faibussowitsch   PetscCall(PetscFree(atb));
20613ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2062baa3c1c6SHong Zhang }
2063baa3c1c6SHong Zhang 
206466976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MatMatTransMult_MPIDense_MPIDense(void *data)
2065d71ae5a4SJacob Faibussowitsch {
20666718818eSStefano Zampini   Mat_MatTransMultDense *abt = (Mat_MatTransMultDense *)data;
2067cc48ffa7SToby Isaac 
2068cc48ffa7SToby Isaac   PetscFunctionBegin;
20699566063dSJacob Faibussowitsch   PetscCall(PetscFree2(abt->buf[0], abt->buf[1]));
20709566063dSJacob Faibussowitsch   PetscCall(PetscFree2(abt->recvcounts, abt->recvdispls));
20719566063dSJacob Faibussowitsch   PetscCall(PetscFree(abt));
20723ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2073cc48ffa7SToby Isaac }
2074cc48ffa7SToby Isaac 
2075d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2076d71ae5a4SJacob Faibussowitsch {
2077baa3c1c6SHong Zhang   Mat_MPIDense          *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
20786718818eSStefano Zampini   Mat_TransMatMultDense *atb;
2079cb20be35SHong Zhang   MPI_Comm               comm;
20806718818eSStefano Zampini   PetscMPIInt            size, *recvcounts;
20816718818eSStefano Zampini   PetscScalar           *carray, *sendbuf;
2082637a0070SStefano Zampini   const PetscScalar     *atbarray;
2083a1f9a5e6SStefano Zampini   PetscInt               i, cN = C->cmap->N, proc, k, j, lda;
2084e68c0b26SHong Zhang   const PetscInt        *ranges;
2085cb20be35SHong Zhang 
2086cb20be35SHong Zhang   PetscFunctionBegin;
20876718818eSStefano Zampini   MatCheckProduct(C, 3);
20885f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
20896718818eSStefano Zampini   atb        = (Mat_TransMatMultDense *)C->product->data;
20906718818eSStefano Zampini   recvcounts = atb->recvcounts;
20916718818eSStefano Zampini   sendbuf    = atb->sendbuf;
20926718818eSStefano Zampini 
20939566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
20949566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
2095e68c0b26SHong Zhang 
2096c5ef1628SHong Zhang   /* compute atbarray = aseq^T * bseq */
2097fb842aefSJose E. Roman   PetscCall(MatTransposeMatMult(a->A, b->A, atb->atb ? MAT_REUSE_MATRIX : MAT_INITIAL_MATRIX, PETSC_DETERMINE, &atb->atb));
2098cb20be35SHong Zhang 
20999566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRanges(C, &ranges));
2100cb20be35SHong Zhang 
2101d016bddeSToby Isaac   if (ranges[1] == C->rmap->N) {
2102d016bddeSToby Isaac     /* all of the values are being reduced to rank 0: optimize this case to use MPI_Reduce and GPU aware MPI if available */
2103d016bddeSToby Isaac     PetscInt           atb_lda, c_lda;
2104d016bddeSToby Isaac     Mat                atb_local = atb->atb;
2105d016bddeSToby Isaac     Mat                atb_alloc = NULL;
2106d016bddeSToby Isaac     Mat                c_local   = c->A;
2107d016bddeSToby Isaac     Mat                c_alloc   = NULL;
2108d016bddeSToby Isaac     PetscMemType       atb_memtype, c_memtype;
2109d016bddeSToby Isaac     const PetscScalar *atb_array = NULL;
2110d016bddeSToby Isaac     MPI_Datatype       vector_type;
2111d016bddeSToby Isaac     PetscScalar       *c_array = NULL;
2112d016bddeSToby Isaac     PetscMPIInt        rank;
2113d016bddeSToby Isaac 
2114d016bddeSToby Isaac     PetscCallMPI(MPI_Comm_rank(comm, &rank));
2115d016bddeSToby Isaac 
2116d016bddeSToby Isaac     PetscCall(MatDenseGetLDA(atb_local, &atb_lda));
2117d016bddeSToby Isaac     if (atb_lda != C->rmap->N) {
2118d016bddeSToby Isaac       // copy atb to a matrix that will have lda == the number of rows
2119d016bddeSToby Isaac       PetscCall(MatDuplicate(atb_local, MAT_DO_NOT_COPY_VALUES, &atb_alloc));
2120d016bddeSToby Isaac       PetscCall(MatCopy(atb_local, atb_alloc, DIFFERENT_NONZERO_PATTERN));
2121d016bddeSToby Isaac       atb_local = atb_alloc;
2122d016bddeSToby Isaac     }
2123d016bddeSToby Isaac 
2124d016bddeSToby Isaac     if (rank == 0) {
2125d016bddeSToby Isaac       PetscCall(MatDenseGetLDA(c_local, &c_lda));
2126d016bddeSToby Isaac       if (c_lda != C->rmap->N) {
2127d016bddeSToby Isaac         // copy c to a matrix that will have lda == the number of rows
2128d016bddeSToby Isaac         PetscCall(MatDuplicate(c_local, MAT_DO_NOT_COPY_VALUES, &c_alloc));
2129d016bddeSToby Isaac         c_local = c_alloc;
2130d016bddeSToby Isaac       }
2131d016bddeSToby Isaac       PetscCall(MatZeroEntries(c_local));
2132d016bddeSToby Isaac     }
2133d016bddeSToby Isaac     /* atb_local and c_local have nrows = lda = A->cmap->N and ncols =
2134d016bddeSToby Isaac      * B->cmap->N: use the a->Mvctx to use the best reduction method */
2135d016bddeSToby Isaac     if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
2136d016bddeSToby Isaac     vector_type = MPIU_SCALAR;
2137d016bddeSToby Isaac     if (B->cmap->N > 1) {
2138b682a77bSToby Isaac       PetscMPIInt mpi_N;
2139b682a77bSToby Isaac 
2140b682a77bSToby Isaac       PetscCall(PetscMPIIntCast(B->cmap->N, &mpi_N));
2141b682a77bSToby Isaac       PetscCallMPI(MPI_Type_contiguous(mpi_N, MPIU_SCALAR, &vector_type));
2142d016bddeSToby Isaac       PetscCallMPI(MPI_Type_commit(&vector_type));
2143d016bddeSToby Isaac     }
2144d016bddeSToby Isaac     PetscCall(MatDenseGetArrayReadAndMemType(atb_local, &atb_array, &atb_memtype));
2145d016bddeSToby Isaac     PetscCall(MatDenseGetArrayWriteAndMemType(c_local, &c_array, &c_memtype));
2146d016bddeSToby Isaac     PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, vector_type, atb_memtype, atb_array, c_memtype, c_array, MPIU_SUM));
2147d016bddeSToby Isaac     PetscCall(PetscSFReduceEnd(a->Mvctx, vector_type, atb_array, c_array, MPIU_SUM));
2148d016bddeSToby Isaac     PetscCall(MatDenseRestoreArrayWriteAndMemType(c_local, &c_array));
2149d016bddeSToby Isaac     PetscCall(MatDenseRestoreArrayReadAndMemType(atb_local, &atb_array));
2150d016bddeSToby Isaac     if (rank == 0 && c_local != c->A) PetscCall(MatCopy(c_local, c->A, DIFFERENT_NONZERO_PATTERN));
2151d016bddeSToby Isaac     if (B->cmap->N > 1) PetscCallMPI(MPI_Type_free(&vector_type));
2152d016bddeSToby Isaac     PetscCall(MatDestroy(&atb_alloc));
2153d016bddeSToby Isaac     PetscCall(MatDestroy(&c_alloc));
2154d016bddeSToby Isaac     PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2155d016bddeSToby Isaac     PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2156d016bddeSToby Isaac     PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2157d016bddeSToby Isaac     PetscFunctionReturn(PETSC_SUCCESS);
2158d016bddeSToby Isaac   }
2159d016bddeSToby Isaac 
2160660d5466SHong Zhang   /* arrange atbarray into sendbuf */
21619566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(atb->atb, &atbarray));
21629566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(atb->atb, &lda));
2163637a0070SStefano Zampini   for (proc = 0, k = 0; proc < size; proc++) {
2164baa3c1c6SHong Zhang     for (j = 0; j < cN; j++) {
2165a1f9a5e6SStefano Zampini       for (i = ranges[proc]; i < ranges[proc + 1]; i++) sendbuf[k++] = atbarray[i + j * lda];
2166cb20be35SHong Zhang     }
2167cb20be35SHong Zhang   }
21689566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(atb->atb, &atbarray));
2169637a0070SStefano Zampini 
2170c5ef1628SHong Zhang   /* sum all atbarray to local values of C */
21719566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(c->A, &carray));
21729566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Reduce_scatter(sendbuf, carray, recvcounts, MPIU_SCALAR, MPIU_SUM, comm));
21739566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(c->A, &carray));
217467af85e8SPierre Jolivet   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
21759566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
21769566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
21773ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2178cb20be35SHong Zhang }
2179cb20be35SHong Zhang 
2180d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2181d71ae5a4SJacob Faibussowitsch {
2182cb20be35SHong Zhang   MPI_Comm               comm;
2183baa3c1c6SHong Zhang   PetscMPIInt            size;
2184660d5466SHong Zhang   PetscInt               cm = A->cmap->n, cM, cN = B->cmap->N;
2185baa3c1c6SHong Zhang   Mat_TransMatMultDense *atb;
218647d993e7Ssuyashtn   PetscBool              cisdense = PETSC_FALSE;
21870acaeb71SStefano Zampini   const PetscInt        *ranges;
2188cb20be35SHong Zhang 
2189cb20be35SHong Zhang   PetscFunctionBegin;
21908f83a4d9SJacob Faibussowitsch   MatCheckProduct(C, 4);
21915f80ce2aSJacob Faibussowitsch   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
21929566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
2193966bd95aSPierre Jolivet   PetscCheck(A->rmap->rstart == B->rmap->rstart && A->rmap->rend == B->rmap->rend, comm, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != B (%" PetscInt_FMT ",%" PetscInt_FMT ")", A->rmap->rstart,
2194966bd95aSPierre Jolivet              A->rmap->rend, B->rmap->rstart, B->rmap->rend);
2195cb20be35SHong Zhang 
21964222ddf1SHong Zhang   /* create matrix product C */
21979566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(C, cm, B->cmap->n, A->cmap->N, B->cmap->N));
219847d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA)
21999566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSECUDA, ""));
220047d993e7Ssuyashtn #endif
220147d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
220247d993e7Ssuyashtn   PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSEHIP, ""));
220347d993e7Ssuyashtn #endif
220448a46eb9SPierre Jolivet   if (!cisdense) PetscCall(MatSetType(C, ((PetscObject)A)->type_name));
22059566063dSJacob Faibussowitsch   PetscCall(MatSetUp(C));
2206baa3c1c6SHong Zhang 
22074222ddf1SHong Zhang   /* create data structure for reuse C */
22089566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
22099566063dSJacob Faibussowitsch   PetscCall(PetscNew(&atb));
22104222ddf1SHong Zhang   cM = C->rmap->N;
22119566063dSJacob Faibussowitsch   PetscCall(PetscMalloc2(cM * cN, &atb->sendbuf, size, &atb->recvcounts));
22129566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRanges(C, &ranges));
22136497c311SBarry Smith   for (PetscMPIInt i = 0; i < size; i++) PetscCall(PetscMPIIntCast((ranges[i + 1] - ranges[i]) * cN, &atb->recvcounts[i]));
22146718818eSStefano Zampini   C->product->data    = atb;
22156718818eSStefano Zampini   C->product->destroy = MatDestroy_MatTransMatMult_MPIDense_MPIDense;
22163ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2217cb20be35SHong Zhang }
2218cb20be35SHong Zhang 
2219d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2220d71ae5a4SJacob Faibussowitsch {
2221cc48ffa7SToby Isaac   MPI_Comm               comm;
2222cc48ffa7SToby Isaac   PetscMPIInt            i, size;
2223cc48ffa7SToby Isaac   PetscInt               maxRows, bufsiz;
2224cc48ffa7SToby Isaac   PetscMPIInt            tag;
22254222ddf1SHong Zhang   PetscInt               alg;
2226cc48ffa7SToby Isaac   Mat_MatTransMultDense *abt;
22274222ddf1SHong Zhang   Mat_Product           *product = C->product;
22284222ddf1SHong Zhang   PetscBool              flg;
2229cc48ffa7SToby Isaac 
2230cc48ffa7SToby Isaac   PetscFunctionBegin;
22316718818eSStefano Zampini   MatCheckProduct(C, 4);
22325f80ce2aSJacob Faibussowitsch   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
22334222ddf1SHong Zhang   /* check local size of A and B */
22345f80ce2aSJacob 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);
2235cc48ffa7SToby Isaac 
22369566063dSJacob Faibussowitsch   PetscCall(PetscStrcmp(product->alg, "allgatherv", &flg));
2237637a0070SStefano Zampini   alg = flg ? 0 : 1;
2238cc48ffa7SToby Isaac 
22394222ddf1SHong Zhang   /* setup matrix product C */
22409566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(C, A->rmap->n, B->rmap->n, A->rmap->N, B->rmap->N));
22419566063dSJacob Faibussowitsch   PetscCall(MatSetType(C, MATMPIDENSE));
22429566063dSJacob Faibussowitsch   PetscCall(MatSetUp(C));
22439566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag));
22444222ddf1SHong Zhang 
22454222ddf1SHong Zhang   /* create data structure for reuse C */
22469566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
22479566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
22489566063dSJacob Faibussowitsch   PetscCall(PetscNew(&abt));
2249cc48ffa7SToby Isaac   abt->tag = tag;
2250faa55883SToby Isaac   abt->alg = alg;
2251faa55883SToby Isaac   switch (alg) {
22524222ddf1SHong Zhang   case 1: /* alg: "cyclic" */
225357508eceSPierre Jolivet     for (maxRows = 0, i = 0; i < size; i++) maxRows = PetscMax(maxRows, B->rmap->range[i + 1] - B->rmap->range[i]);
2254cc48ffa7SToby Isaac     bufsiz = A->cmap->N * maxRows;
2255f4f49eeaSPierre Jolivet     PetscCall(PetscMalloc2(bufsiz, &abt->buf[0], bufsiz, &abt->buf[1]));
2256faa55883SToby Isaac     break;
22574222ddf1SHong Zhang   default: /* alg: "allgatherv" */
2258f4f49eeaSPierre Jolivet     PetscCall(PetscMalloc2(B->rmap->n * B->cmap->N, &abt->buf[0], B->rmap->N * B->cmap->N, &abt->buf[1]));
2259f4f49eeaSPierre Jolivet     PetscCall(PetscMalloc2(size, &abt->recvcounts, size + 1, &abt->recvdispls));
22606497c311SBarry Smith     for (i = 0; i <= size; i++) PetscCall(PetscMPIIntCast(B->rmap->range[i] * A->cmap->N, &abt->recvdispls[i]));
22616497c311SBarry Smith     for (i = 0; i < size; i++) PetscCall(PetscMPIIntCast(abt->recvdispls[i + 1] - abt->recvdispls[i], &abt->recvcounts[i]));
2262faa55883SToby Isaac     break;
2263faa55883SToby Isaac   }
2264cc48ffa7SToby Isaac 
22656718818eSStefano Zampini   C->product->data                = abt;
22666718818eSStefano Zampini   C->product->destroy             = MatDestroy_MatMatTransMult_MPIDense_MPIDense;
22674222ddf1SHong Zhang   C->ops->mattransposemultnumeric = MatMatTransposeMultNumeric_MPIDense_MPIDense;
22683ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2269cc48ffa7SToby Isaac }
2270cc48ffa7SToby Isaac 
2271d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(Mat A, Mat B, Mat C)
2272d71ae5a4SJacob Faibussowitsch {
2273cc48ffa7SToby Isaac   Mat_MPIDense          *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
22746718818eSStefano Zampini   Mat_MatTransMultDense *abt;
2275cc48ffa7SToby Isaac   MPI_Comm               comm;
22766497c311SBarry Smith   PetscMPIInt            rank, size, sendto, recvfrom, recvisfrom;
2277f4259b30SLisandro Dalcin   PetscScalar           *sendbuf, *recvbuf = NULL, *cv;
22786497c311SBarry Smith   PetscInt               i, cK             = A->cmap->N, sendsiz, recvsiz, k, j, bn;
2279cc48ffa7SToby Isaac   PetscScalar            _DOne = 1.0, _DZero = 0.0;
2280637a0070SStefano Zampini   const PetscScalar     *av, *bv;
228123fff9afSBarry Smith   PetscBLASInt           cm, cn, ck, alda, blda = 0, clda;
2282cc48ffa7SToby Isaac   MPI_Request            reqs[2];
2283cc48ffa7SToby Isaac   const PetscInt        *ranges;
2284cc48ffa7SToby Isaac 
2285cc48ffa7SToby Isaac   PetscFunctionBegin;
22866718818eSStefano Zampini   MatCheckProduct(C, 3);
22875f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
22886718818eSStefano Zampini   abt = (Mat_MatTransMultDense *)C->product->data;
22899566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
22909566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(comm, &rank));
22919566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
22929566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, &av));
22939566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(b->A, &bv));
22949566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(c->A, &cv));
22959566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &i));
22969566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &alda));
22979566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(b->A, &i));
22989566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &blda));
22999566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(c->A, &i));
23009566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &clda));
23019566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRanges(B, &ranges));
2302cc48ffa7SToby Isaac   bn = B->rmap->n;
2303637a0070SStefano Zampini   if (blda == bn) {
2304637a0070SStefano Zampini     sendbuf = (PetscScalar *)bv;
2305cc48ffa7SToby Isaac   } else {
2306cc48ffa7SToby Isaac     sendbuf = abt->buf[0];
2307cc48ffa7SToby Isaac     for (k = 0, i = 0; i < cK; i++) {
2308ad540459SPierre Jolivet       for (j = 0; j < bn; j++, k++) sendbuf[k] = bv[i * blda + j];
2309cc48ffa7SToby Isaac     }
2310cc48ffa7SToby Isaac   }
2311cc48ffa7SToby Isaac   if (size > 1) {
2312cc48ffa7SToby Isaac     sendto   = (rank + size - 1) % size;
2313cc48ffa7SToby Isaac     recvfrom = (rank + size + 1) % size;
2314cc48ffa7SToby Isaac   } else {
2315cc48ffa7SToby Isaac     sendto = recvfrom = 0;
2316cc48ffa7SToby Isaac   }
23179566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(cK, &ck));
23189566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm));
2319cc48ffa7SToby Isaac   recvisfrom = rank;
2320cc48ffa7SToby Isaac   for (i = 0; i < size; i++) {
2321cc48ffa7SToby Isaac     /* we have finished receiving in sending, bufs can be read/modified */
23226497c311SBarry Smith     PetscMPIInt nextrecvisfrom = (recvisfrom + 1) % size; /* which process the next recvbuf will originate on */
2323cc48ffa7SToby Isaac     PetscInt    nextbn         = ranges[nextrecvisfrom + 1] - ranges[nextrecvisfrom];
2324cc48ffa7SToby Isaac 
2325cc48ffa7SToby Isaac     if (nextrecvisfrom != rank) {
2326cc48ffa7SToby Isaac       /* start the cyclic sends from sendbuf, to recvbuf (which will switch to sendbuf) */
2327cc48ffa7SToby Isaac       sendsiz = cK * bn;
2328cc48ffa7SToby Isaac       recvsiz = cK * nextbn;
2329cc48ffa7SToby Isaac       recvbuf = (i & 1) ? abt->buf[0] : abt->buf[1];
23306497c311SBarry Smith       PetscCallMPI(MPIU_Isend(sendbuf, sendsiz, MPIU_SCALAR, sendto, abt->tag, comm, &reqs[0]));
23316497c311SBarry Smith       PetscCallMPI(MPIU_Irecv(recvbuf, recvsiz, MPIU_SCALAR, recvfrom, abt->tag, comm, &reqs[1]));
2332cc48ffa7SToby Isaac     }
2333cc48ffa7SToby Isaac 
2334cc48ffa7SToby Isaac     /* local aseq * sendbuf^T */
23359566063dSJacob Faibussowitsch     PetscCall(PetscBLASIntCast(ranges[recvisfrom + 1] - ranges[recvisfrom], &cn));
2336792fecdfSBarry Smith     if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "T", &cm, &cn, &ck, &_DOne, av, &alda, sendbuf, &cn, &_DZero, cv + clda * ranges[recvisfrom], &clda));
2337cc48ffa7SToby Isaac 
2338cc48ffa7SToby Isaac     if (nextrecvisfrom != rank) {
2339cc48ffa7SToby Isaac       /* wait for the sends and receives to complete, swap sendbuf and recvbuf */
23409566063dSJacob Faibussowitsch       PetscCallMPI(MPI_Waitall(2, reqs, MPI_STATUSES_IGNORE));
2341cc48ffa7SToby Isaac     }
2342cc48ffa7SToby Isaac     bn         = nextbn;
2343cc48ffa7SToby Isaac     recvisfrom = nextrecvisfrom;
2344cc48ffa7SToby Isaac     sendbuf    = recvbuf;
2345cc48ffa7SToby Isaac   }
23469566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, &av));
23479566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
23489566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(c->A, &cv));
234967af85e8SPierre Jolivet   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
23509566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
23519566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
23523ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2353cc48ffa7SToby Isaac }
2354cc48ffa7SToby Isaac 
2355d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(Mat A, Mat B, Mat C)
2356d71ae5a4SJacob Faibussowitsch {
2357faa55883SToby Isaac   Mat_MPIDense          *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
23586718818eSStefano Zampini   Mat_MatTransMultDense *abt;
2359faa55883SToby Isaac   MPI_Comm               comm;
23606497c311SBarry Smith   PetscMPIInt            size, ibn;
2361637a0070SStefano Zampini   PetscScalar           *cv, *sendbuf, *recvbuf;
2362637a0070SStefano Zampini   const PetscScalar     *av, *bv;
2363637a0070SStefano Zampini   PetscInt               blda, i, cK = A->cmap->N, k, j, bn;
2364faa55883SToby Isaac   PetscScalar            _DOne = 1.0, _DZero = 0.0;
2365637a0070SStefano Zampini   PetscBLASInt           cm, cn, ck, alda, clda;
2366faa55883SToby Isaac 
2367faa55883SToby Isaac   PetscFunctionBegin;
23686718818eSStefano Zampini   MatCheckProduct(C, 3);
23695f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
23706718818eSStefano Zampini   abt = (Mat_MatTransMultDense *)C->product->data;
23719566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
23729566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
23739566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, &av));
23749566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(b->A, &bv));
23759566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(c->A, &cv));
23769566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &i));
23779566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &alda));
23789566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(b->A, &blda));
23799566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(c->A, &i));
23809566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &clda));
2381faa55883SToby Isaac   /* copy transpose of B into buf[0] */
2382faa55883SToby Isaac   bn      = B->rmap->n;
2383faa55883SToby Isaac   sendbuf = abt->buf[0];
2384faa55883SToby Isaac   recvbuf = abt->buf[1];
2385faa55883SToby Isaac   for (k = 0, j = 0; j < bn; j++) {
2386ad540459SPierre Jolivet     for (i = 0; i < cK; i++, k++) sendbuf[k] = bv[i * blda + j];
2387faa55883SToby Isaac   }
23889566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
23896497c311SBarry Smith   PetscCall(PetscMPIIntCast(bn * cK, &ibn));
23906497c311SBarry Smith   PetscCallMPI(MPI_Allgatherv(sendbuf, ibn, MPIU_SCALAR, recvbuf, abt->recvcounts, abt->recvdispls, MPIU_SCALAR, comm));
23919566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(cK, &ck));
23929566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm));
23939566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(c->A->cmap->n, &cn));
2394792fecdfSBarry Smith   if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "N", &cm, &cn, &ck, &_DOne, av, &alda, recvbuf, &ck, &_DZero, cv, &clda));
23959566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, &av));
23969566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
23979566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(c->A, &cv));
239867af85e8SPierre Jolivet   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
23999566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
24009566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
24013ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2402faa55883SToby Isaac }
2403faa55883SToby Isaac 
2404d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2405d71ae5a4SJacob Faibussowitsch {
24066718818eSStefano Zampini   Mat_MatTransMultDense *abt;
2407faa55883SToby Isaac 
2408faa55883SToby Isaac   PetscFunctionBegin;
24096718818eSStefano Zampini   MatCheckProduct(C, 3);
24105f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
24116718818eSStefano Zampini   abt = (Mat_MatTransMultDense *)C->product->data;
2412faa55883SToby Isaac   switch (abt->alg) {
2413d71ae5a4SJacob Faibussowitsch   case 1:
2414d71ae5a4SJacob Faibussowitsch     PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(A, B, C));
2415d71ae5a4SJacob Faibussowitsch     break;
2416d71ae5a4SJacob Faibussowitsch   default:
2417d71ae5a4SJacob Faibussowitsch     PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(A, B, C));
2418d71ae5a4SJacob Faibussowitsch     break;
2419faa55883SToby Isaac   }
24203ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2421faa55883SToby Isaac }
2422faa55883SToby Isaac 
242366976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MatMatMult_MPIDense_MPIDense(void *data)
2424d71ae5a4SJacob Faibussowitsch {
24256718818eSStefano Zampini   Mat_MatMultDense *ab = (Mat_MatMultDense *)data;
2426320f2790SHong Zhang 
2427320f2790SHong Zhang   PetscFunctionBegin;
24289566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&ab->Ce));
24299566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&ab->Ae));
24309566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&ab->Be));
24319566063dSJacob Faibussowitsch   PetscCall(PetscFree(ab));
24323ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2433320f2790SHong Zhang }
2434320f2790SHong Zhang 
24355d8c7819SPierre Jolivet static PetscErrorCode MatMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2436d71ae5a4SJacob Faibussowitsch {
24376718818eSStefano Zampini   Mat_MatMultDense *ab;
24385d8c7819SPierre Jolivet   Mat_MPIDense     *mdn = (Mat_MPIDense *)A->data;
2439d016bddeSToby Isaac   Mat_MPIDense     *b   = (Mat_MPIDense *)B->data;
2440320f2790SHong Zhang 
2441320f2790SHong Zhang   PetscFunctionBegin;
24426718818eSStefano Zampini   MatCheckProduct(C, 3);
24435f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Missing product data");
24446718818eSStefano Zampini   ab = (Mat_MatMultDense *)C->product->data;
24455d8c7819SPierre Jolivet   if (ab->Ae && ab->Ce) {
24465d8c7819SPierre Jolivet #if PetscDefined(HAVE_ELEMENTAL)
24479566063dSJacob Faibussowitsch     PetscCall(MatConvert_MPIDense_Elemental(A, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Ae));
24489566063dSJacob Faibussowitsch     PetscCall(MatConvert_MPIDense_Elemental(B, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Be));
24499566063dSJacob Faibussowitsch     PetscCall(MatMatMultNumeric_Elemental(ab->Ae, ab->Be, ab->Ce));
24509566063dSJacob Faibussowitsch     PetscCall(MatConvert(ab->Ce, MATMPIDENSE, MAT_REUSE_MATRIX, &C));
24515d8c7819SPierre Jolivet #else
24525d8c7819SPierre Jolivet     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined");
24535d8c7819SPierre Jolivet #endif
24545d8c7819SPierre Jolivet   } else {
2455d016bddeSToby Isaac     MPI_Comm           comm;
24565d8c7819SPierre Jolivet     const PetscScalar *read;
24575d8c7819SPierre Jolivet     PetscScalar       *write;
24585d8c7819SPierre Jolivet     PetscInt           lda;
2459d016bddeSToby Isaac     const PetscInt    *ranges;
2460d016bddeSToby Isaac     PetscMPIInt        size;
24615d8c7819SPierre Jolivet 
2462d016bddeSToby Isaac     if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A)); /* cannot be done during the symbolic phase because of possible calls to MatProductReplaceMats() */
2463d016bddeSToby Isaac     comm = PetscObjectComm((PetscObject)B);
2464d016bddeSToby Isaac     PetscCallMPI(MPI_Comm_size(comm, &size));
2465d016bddeSToby Isaac     PetscCall(PetscLayoutGetRanges(B->rmap, &ranges));
2466d016bddeSToby Isaac     if (ranges[1] == ranges[size]) {
2467d016bddeSToby Isaac       // optimize for the case where the B matrix is broadcast from rank 0
2468d016bddeSToby Isaac       PetscInt           b_lda, be_lda;
2469d016bddeSToby Isaac       Mat                b_local  = b->A;
2470d016bddeSToby Isaac       Mat                b_alloc  = NULL;
2471d016bddeSToby Isaac       Mat                be_local = ab->Be;
2472d016bddeSToby Isaac       Mat                be_alloc = NULL;
2473d016bddeSToby Isaac       PetscMemType       b_memtype, be_memtype;
2474d016bddeSToby Isaac       const PetscScalar *b_array = NULL;
2475d016bddeSToby Isaac       MPI_Datatype       vector_type;
2476d016bddeSToby Isaac       PetscScalar       *be_array = NULL;
2477d016bddeSToby Isaac       PetscMPIInt        rank;
2478d016bddeSToby Isaac 
2479d016bddeSToby Isaac       PetscCallMPI(MPI_Comm_rank(comm, &rank));
2480d016bddeSToby Isaac       PetscCall(MatDenseGetLDA(be_local, &be_lda));
2481d016bddeSToby Isaac       if (be_lda != B->rmap->N) {
2482d016bddeSToby Isaac         PetscCall(MatDuplicate(be_local, MAT_DO_NOT_COPY_VALUES, &be_alloc));
2483d016bddeSToby Isaac         be_local = be_alloc;
2484d016bddeSToby Isaac       }
2485d016bddeSToby Isaac 
2486d016bddeSToby Isaac       if (rank == 0) {
2487d016bddeSToby Isaac         PetscCall(MatDenseGetLDA(b_local, &b_lda));
2488d016bddeSToby Isaac         if (b_lda != B->rmap->N) {
2489d016bddeSToby Isaac           PetscCall(MatDuplicate(b_local, MAT_DO_NOT_COPY_VALUES, &b_alloc));
2490d016bddeSToby Isaac           PetscCall(MatCopy(b_local, b_alloc, DIFFERENT_NONZERO_PATTERN));
2491d016bddeSToby Isaac           b_local = b_alloc;
2492d016bddeSToby Isaac         }
2493d016bddeSToby Isaac       }
2494b682a77bSToby Isaac       vector_type = MPIU_SCALAR;
2495b682a77bSToby Isaac       if (B->cmap->N > 1) {
2496b682a77bSToby Isaac         PetscMPIInt mpi_N;
2497b682a77bSToby Isaac 
2498b682a77bSToby Isaac         PetscCall(PetscMPIIntCast(B->cmap->N, &mpi_N));
2499b682a77bSToby Isaac         PetscCallMPI(MPI_Type_contiguous(mpi_N, MPIU_SCALAR, &vector_type));
2500d016bddeSToby Isaac         PetscCallMPI(MPI_Type_commit(&vector_type));
2501b682a77bSToby Isaac       }
2502d016bddeSToby Isaac       PetscCall(MatDenseGetArrayReadAndMemType(b_local, &b_array, &b_memtype));
2503d016bddeSToby Isaac       PetscCall(MatDenseGetArrayWriteAndMemType(be_local, &be_array, &be_memtype));
2504d016bddeSToby Isaac       PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, vector_type, b_memtype, b_array, be_memtype, be_array, MPI_REPLACE));
2505d016bddeSToby Isaac       PetscCall(PetscSFBcastEnd(mdn->Mvctx, vector_type, b_array, be_array, MPI_REPLACE));
2506d016bddeSToby Isaac       PetscCall(MatDenseRestoreArrayWriteAndMemType(be_local, &be_array));
2507d016bddeSToby Isaac       PetscCall(MatDenseRestoreArrayReadAndMemType(b_local, &b_array));
2508d016bddeSToby Isaac       if (be_local != ab->Be) PetscCall(MatCopy(be_local, ab->Be, DIFFERENT_NONZERO_PATTERN));
2509b682a77bSToby Isaac       if (B->cmap->N > 1) PetscCallMPI(MPI_Type_free(&vector_type));
2510d016bddeSToby Isaac       PetscCall(MatDestroy(&be_alloc));
2511d016bddeSToby Isaac       PetscCall(MatDestroy(&b_alloc));
2512d016bddeSToby Isaac     } else {
25135d8c7819SPierre Jolivet       PetscCall(MatDenseGetLDA(B, &lda));
25145d8c7819SPierre Jolivet       PetscCall(MatDenseGetArrayRead(B, &read));
25155d8c7819SPierre Jolivet       PetscCall(MatDenseGetArrayWrite(ab->Be, &write));
25165d8c7819SPierre Jolivet       for (PetscInt i = 0; i < C->cmap->N; ++i) {
25175d8c7819SPierre Jolivet         PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE));
25185d8c7819SPierre Jolivet         PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE));
25195d8c7819SPierre Jolivet       }
25205d8c7819SPierre Jolivet       PetscCall(MatDenseRestoreArrayWrite(ab->Be, &write));
25215d8c7819SPierre Jolivet       PetscCall(MatDenseRestoreArrayRead(B, &read));
2522d016bddeSToby Isaac     }
2523f4f49eeaSPierre Jolivet     PetscCall(MatMatMultNumeric_SeqDense_SeqDense(((Mat_MPIDense *)A->data)->A, ab->Be, ((Mat_MPIDense *)C->data)->A));
25245d8c7819SPierre Jolivet   }
25253ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2526320f2790SHong Zhang }
2527320f2790SHong Zhang 
2528d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2529d71ae5a4SJacob Faibussowitsch {
25305d8c7819SPierre Jolivet   Mat_Product      *product = C->product;
25315d8c7819SPierre Jolivet   PetscInt          alg;
2532320f2790SHong Zhang   Mat_MatMultDense *ab;
25335d8c7819SPierre Jolivet   PetscBool         flg;
2534320f2790SHong Zhang 
2535320f2790SHong Zhang   PetscFunctionBegin;
25366718818eSStefano Zampini   MatCheckProduct(C, 4);
25375f80ce2aSJacob Faibussowitsch   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
25384222ddf1SHong Zhang   /* check local size of A and B */
25395d8c7819SPierre 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 ")",
25405d8c7819SPierre Jolivet              A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
2541320f2790SHong Zhang 
25425d8c7819SPierre Jolivet   PetscCall(PetscStrcmp(product->alg, "petsc", &flg));
25435d8c7819SPierre Jolivet   alg = flg ? 0 : 1;
25444222ddf1SHong Zhang 
25454222ddf1SHong Zhang   /* setup C */
25465d8c7819SPierre Jolivet   PetscCall(MatSetSizes(C, A->rmap->n, B->cmap->n, A->rmap->N, B->cmap->N));
25475d8c7819SPierre Jolivet   PetscCall(MatSetType(C, MATMPIDENSE));
25489566063dSJacob Faibussowitsch   PetscCall(MatSetUp(C));
2549320f2790SHong Zhang 
2550320f2790SHong Zhang   /* create data structure for reuse Cdense */
25519566063dSJacob Faibussowitsch   PetscCall(PetscNew(&ab));
25525d8c7819SPierre Jolivet 
25535d8c7819SPierre Jolivet   switch (alg) {
25545d8c7819SPierre Jolivet   case 1: /* alg: "elemental" */
25555d8c7819SPierre Jolivet #if PetscDefined(HAVE_ELEMENTAL)
25565d8c7819SPierre Jolivet     /* create elemental matrices Ae and Be */
25575d8c7819SPierre Jolivet     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &ab->Ae));
25585d8c7819SPierre Jolivet     PetscCall(MatSetSizes(ab->Ae, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N));
25595d8c7819SPierre Jolivet     PetscCall(MatSetType(ab->Ae, MATELEMENTAL));
25605d8c7819SPierre Jolivet     PetscCall(MatSetUp(ab->Ae));
25615d8c7819SPierre Jolivet     PetscCall(MatSetOption(ab->Ae, MAT_ROW_ORIENTED, PETSC_FALSE));
25625d8c7819SPierre Jolivet 
25635d8c7819SPierre Jolivet     PetscCall(MatCreate(PetscObjectComm((PetscObject)B), &ab->Be));
25645d8c7819SPierre Jolivet     PetscCall(MatSetSizes(ab->Be, PETSC_DECIDE, PETSC_DECIDE, B->rmap->N, B->cmap->N));
25655d8c7819SPierre Jolivet     PetscCall(MatSetType(ab->Be, MATELEMENTAL));
25665d8c7819SPierre Jolivet     PetscCall(MatSetUp(ab->Be));
25675d8c7819SPierre Jolivet     PetscCall(MatSetOption(ab->Be, MAT_ROW_ORIENTED, PETSC_FALSE));
25685d8c7819SPierre Jolivet 
25695d8c7819SPierre Jolivet     /* compute symbolic Ce = Ae*Be */
25705d8c7819SPierre Jolivet     PetscCall(MatCreate(PetscObjectComm((PetscObject)C), &ab->Ce));
25715d8c7819SPierre Jolivet     PetscCall(MatMatMultSymbolic_Elemental(ab->Ae, ab->Be, fill, ab->Ce));
25725d8c7819SPierre Jolivet #else
25735d8c7819SPierre Jolivet     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined");
25745d8c7819SPierre Jolivet #endif
25755d8c7819SPierre Jolivet     break;
25765d8c7819SPierre Jolivet   default: /* alg: "petsc" */
25775d8c7819SPierre Jolivet     ab->Ae = NULL;
25785d8c7819SPierre Jolivet     PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, A->cmap->N, B->cmap->N, NULL, &ab->Be));
25795d8c7819SPierre Jolivet     ab->Ce = NULL;
25805d8c7819SPierre Jolivet     break;
25815d8c7819SPierre Jolivet   }
25826718818eSStefano Zampini 
25836718818eSStefano Zampini   C->product->data       = ab;
25846718818eSStefano Zampini   C->product->destroy    = MatDestroy_MatMatMult_MPIDense_MPIDense;
25854222ddf1SHong Zhang   C->ops->matmultnumeric = MatMatMultNumeric_MPIDense_MPIDense;
25863ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2587320f2790SHong Zhang }
25882ef1f0ffSBarry Smith 
2589d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_AB(Mat C)
2590d71ae5a4SJacob Faibussowitsch {
25915d8c7819SPierre Jolivet   Mat_Product *product     = C->product;
25925d8c7819SPierre Jolivet   const char  *algTypes[2] = {"petsc", "elemental"};
25935d8c7819SPierre Jolivet   PetscInt     alg, nalg = PetscDefined(HAVE_ELEMENTAL) ? 2 : 1;
25945d8c7819SPierre Jolivet   PetscBool    flg = PETSC_FALSE;
25955d8c7819SPierre Jolivet 
25964d86920dSPierre Jolivet   PetscFunctionBegin;
25975d8c7819SPierre Jolivet   /* Set default algorithm */
2598f0b74427SPierre Jolivet   alg = 0; /* default is PETSc */
25995d8c7819SPierre Jolivet   PetscCall(PetscStrcmp(product->alg, "default", &flg));
2600835f2295SStefano Zampini   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
26015d8c7819SPierre Jolivet 
26025d8c7819SPierre Jolivet   /* Get runtime option */
26035d8c7819SPierre Jolivet   PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_AB", "Mat");
26045d8c7819SPierre Jolivet   PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_AB", algTypes, nalg, algTypes[alg], &alg, &flg));
26055d8c7819SPierre Jolivet   PetscOptionsEnd();
2606835f2295SStefano Zampini   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
26075d8c7819SPierre Jolivet 
26084222ddf1SHong Zhang   C->ops->matmultsymbolic = MatMatMultSymbolic_MPIDense_MPIDense;
26094222ddf1SHong Zhang   C->ops->productsymbolic = MatProductSymbolic_AB;
26103ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2611320f2790SHong Zhang }
261286aefd0dSHong Zhang 
2613d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_AtB(Mat C)
2614d71ae5a4SJacob Faibussowitsch {
26154222ddf1SHong Zhang   Mat_Product *product = C->product;
26164222ddf1SHong Zhang   Mat          A = product->A, B = product->B;
26174222ddf1SHong Zhang 
26184222ddf1SHong Zhang   PetscFunctionBegin;
261907e90273SPierre 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 ")",
262007e90273SPierre Jolivet              A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
26216718818eSStefano Zampini   C->ops->transposematmultsymbolic = MatTransposeMatMultSymbolic_MPIDense_MPIDense;
26226718818eSStefano Zampini   C->ops->productsymbolic          = MatProductSymbolic_AtB;
26233ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
26244222ddf1SHong Zhang }
26254222ddf1SHong Zhang 
2626d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_ABt(Mat C)
2627d71ae5a4SJacob Faibussowitsch {
26284222ddf1SHong Zhang   Mat_Product *product     = C->product;
26294222ddf1SHong Zhang   const char  *algTypes[2] = {"allgatherv", "cyclic"};
26304222ddf1SHong Zhang   PetscInt     alg, nalg = 2;
26314222ddf1SHong Zhang   PetscBool    flg = PETSC_FALSE;
26324222ddf1SHong Zhang 
26334222ddf1SHong Zhang   PetscFunctionBegin;
26344222ddf1SHong Zhang   /* Set default algorithm */
26354222ddf1SHong Zhang   alg = 0; /* default is allgatherv */
26369566063dSJacob Faibussowitsch   PetscCall(PetscStrcmp(product->alg, "default", &flg));
2637835f2295SStefano Zampini   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
26384222ddf1SHong Zhang 
26394222ddf1SHong Zhang   /* Get runtime option */
26404222ddf1SHong Zhang   if (product->api_user) {
2641d0609cedSBarry Smith     PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatMatTransposeMult", "Mat");
26429566063dSJacob Faibussowitsch     PetscCall(PetscOptionsEList("-matmattransmult_mpidense_mpidense_via", "Algorithmic approach", "MatMatTransposeMult", algTypes, nalg, algTypes[alg], &alg, &flg));
2643d0609cedSBarry Smith     PetscOptionsEnd();
26444222ddf1SHong Zhang   } else {
2645d0609cedSBarry Smith     PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_ABt", "Mat");
26469566063dSJacob Faibussowitsch     PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_ABt", algTypes, nalg, algTypes[alg], &alg, &flg));
2647d0609cedSBarry Smith     PetscOptionsEnd();
26484222ddf1SHong Zhang   }
2649835f2295SStefano Zampini   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
26504222ddf1SHong Zhang 
26514222ddf1SHong Zhang   C->ops->mattransposemultsymbolic = MatMatTransposeMultSymbolic_MPIDense_MPIDense;
26524222ddf1SHong Zhang   C->ops->productsymbolic          = MatProductSymbolic_ABt;
26533ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
26544222ddf1SHong Zhang }
26554222ddf1SHong Zhang 
26565d8c7819SPierre Jolivet static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat C)
2657d71ae5a4SJacob Faibussowitsch {
26584222ddf1SHong Zhang   Mat_Product *product = C->product;
26594222ddf1SHong Zhang 
26604222ddf1SHong Zhang   PetscFunctionBegin;
26614222ddf1SHong Zhang   switch (product->type) {
2662d71ae5a4SJacob Faibussowitsch   case MATPRODUCT_AB:
2663d71ae5a4SJacob Faibussowitsch     PetscCall(MatProductSetFromOptions_MPIDense_AB(C));
2664d71ae5a4SJacob Faibussowitsch     break;
2665d71ae5a4SJacob Faibussowitsch   case MATPRODUCT_AtB:
2666d71ae5a4SJacob Faibussowitsch     PetscCall(MatProductSetFromOptions_MPIDense_AtB(C));
2667d71ae5a4SJacob Faibussowitsch     break;
2668d71ae5a4SJacob Faibussowitsch   case MATPRODUCT_ABt:
2669d71ae5a4SJacob Faibussowitsch     PetscCall(MatProductSetFromOptions_MPIDense_ABt(C));
2670d71ae5a4SJacob Faibussowitsch     break;
2671d71ae5a4SJacob Faibussowitsch   default:
2672d71ae5a4SJacob Faibussowitsch     break;
26734222ddf1SHong Zhang   }
26743ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
26754222ddf1SHong Zhang }
2676