xref: /petsc/src/mat/impls/dense/mpi/mpidense.c (revision b682a77b1328162988dce990bc8fc42eb4db6521)
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");
175b49de8d1SLois Curfman McInnes     if (idxm[i] >= rstart && idxm[i] < rend) {
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       }
182e7e72b3dSBarry Smith     } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "Only local values currently supported");
1838965ea79SLois Curfman McInnes   }
1843ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1858965ea79SLois Curfman McInnes }
1868965ea79SLois Curfman McInnes 
187d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetLDA_MPIDense(Mat A, PetscInt *lda)
188d71ae5a4SJacob Faibussowitsch {
18949a6ff4bSBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
19049a6ff4bSBarry Smith 
19149a6ff4bSBarry Smith   PetscFunctionBegin;
1929566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, lda));
1933ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
19449a6ff4bSBarry Smith }
19549a6ff4bSBarry Smith 
196d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseSetLDA_MPIDense(Mat A, PetscInt lda)
197d71ae5a4SJacob Faibussowitsch {
198ad16ce7aSStefano Zampini   Mat_MPIDense *a     = (Mat_MPIDense *)A->data;
19947d993e7Ssuyashtn   MatType       mtype = MATSEQDENSE;
200ad16ce7aSStefano Zampini 
201ad16ce7aSStefano Zampini   PetscFunctionBegin;
20217359960SJose E. Roman   if (!a->A) {
20328b400f6SJacob Faibussowitsch     PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2049566063dSJacob Faibussowitsch     PetscCall(PetscLayoutSetUp(A->rmap));
2059566063dSJacob Faibussowitsch     PetscCall(PetscLayoutSetUp(A->cmap));
2069566063dSJacob Faibussowitsch     PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
2079566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(a->A, A->rmap->n, A->cmap->N, A->rmap->n, A->cmap->N));
208a51cd166SPierre Jolivet #if PetscDefined(HAVE_CUDA)
20947d993e7Ssuyashtn     PetscBool iscuda;
2109566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSECUDA, &iscuda));
21147d993e7Ssuyashtn     if (iscuda) mtype = MATSEQDENSECUDA;
212a51cd166SPierre Jolivet #elif PetscDefined(HAVE_HIP)
21347d993e7Ssuyashtn     PetscBool iship;
21447d993e7Ssuyashtn     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSEHIP, &iship));
21547d993e7Ssuyashtn     if (iship) mtype = MATSEQDENSEHIP;
21647d993e7Ssuyashtn #endif
21747d993e7Ssuyashtn     PetscCall(MatSetType(a->A, mtype));
21817359960SJose E. Roman   }
2199566063dSJacob Faibussowitsch   PetscCall(MatDenseSetLDA(a->A, lda));
2203ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
221ad16ce7aSStefano Zampini }
222ad16ce7aSStefano Zampini 
223d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetArray_MPIDense(Mat A, PetscScalar **array)
224d71ae5a4SJacob Faibussowitsch {
225ff14e315SSatish Balay   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
226ff14e315SSatish Balay 
2273a40ed3dSBarry Smith   PetscFunctionBegin;
22828b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2299566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(a->A, array));
2303ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
231ff14e315SSatish Balay }
232ff14e315SSatish Balay 
233f1e99121SPierre Jolivet static PetscErrorCode MatDenseGetArrayRead_MPIDense(Mat A, PetscScalar **array)
234d71ae5a4SJacob Faibussowitsch {
2358572280aSBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
2368572280aSBarry Smith 
2378572280aSBarry Smith   PetscFunctionBegin;
23828b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
239f1e99121SPierre Jolivet   PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)array));
2403ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2418572280aSBarry Smith }
2428572280aSBarry Smith 
243d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetArrayWrite_MPIDense(Mat A, PetscScalar **array)
244d71ae5a4SJacob Faibussowitsch {
2456947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
2466947451fSStefano Zampini 
2476947451fSStefano Zampini   PetscFunctionBegin;
24828b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2499566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(a->A, array));
2503ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2516947451fSStefano Zampini }
2526947451fSStefano Zampini 
253d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDensePlaceArray_MPIDense(Mat A, const PetscScalar *array)
254d71ae5a4SJacob Faibussowitsch {
255d3042a70SBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
256d3042a70SBarry Smith 
257d3042a70SBarry Smith   PetscFunctionBegin;
25828b400f6SJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
25928b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2609566063dSJacob Faibussowitsch   PetscCall(MatDensePlaceArray(a->A, array));
2613ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
262d3042a70SBarry Smith }
263d3042a70SBarry Smith 
264d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseResetArray_MPIDense(Mat A)
265d71ae5a4SJacob Faibussowitsch {
266d3042a70SBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
267d3042a70SBarry Smith 
268d3042a70SBarry Smith   PetscFunctionBegin;
26928b400f6SJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
27028b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2719566063dSJacob Faibussowitsch   PetscCall(MatDenseResetArray(a->A));
2723ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
273d3042a70SBarry Smith }
274d3042a70SBarry Smith 
275d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseReplaceArray_MPIDense(Mat A, const PetscScalar *array)
276d71ae5a4SJacob Faibussowitsch {
277d5ea218eSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
278d5ea218eSStefano Zampini 
279d5ea218eSStefano Zampini   PetscFunctionBegin;
28028b400f6SJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
28128b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2829566063dSJacob Faibussowitsch   PetscCall(MatDenseReplaceArray(a->A, array));
2833ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
284d5ea218eSStefano Zampini }
285d5ea218eSStefano Zampini 
286d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatCreateSubMatrix_MPIDense(Mat A, IS isrow, IS iscol, MatReuse scall, Mat *B)
287d71ae5a4SJacob Faibussowitsch {
288ca3fa75bSLois Curfman McInnes   Mat_MPIDense      *mat = (Mat_MPIDense *)A->data, *newmatd;
289637a0070SStefano Zampini   PetscInt           lda, i, j, rstart, rend, nrows, ncols, Ncols, nlrows, nlcols;
2905d0c19d7SBarry Smith   const PetscInt    *irow, *icol;
291637a0070SStefano Zampini   const PetscScalar *v;
292637a0070SStefano Zampini   PetscScalar       *bv;
293ca3fa75bSLois Curfman McInnes   Mat                newmat;
2944aa3045dSJed Brown   IS                 iscol_local;
29542a884f0SBarry Smith   MPI_Comm           comm_is, comm_mat;
296ca3fa75bSLois Curfman McInnes 
297ca3fa75bSLois Curfman McInnes   PetscFunctionBegin;
2989566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm_mat));
2999566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)iscol, &comm_is));
30008401ef6SPierre Jolivet   PetscCheck(comm_mat == comm_is, PETSC_COMM_SELF, PETSC_ERR_ARG_NOTSAMECOMM, "IS communicator must match matrix communicator");
30142a884f0SBarry Smith 
3029566063dSJacob Faibussowitsch   PetscCall(ISAllGather(iscol, &iscol_local));
3039566063dSJacob Faibussowitsch   PetscCall(ISGetIndices(isrow, &irow));
3049566063dSJacob Faibussowitsch   PetscCall(ISGetIndices(iscol_local, &icol));
3059566063dSJacob Faibussowitsch   PetscCall(ISGetLocalSize(isrow, &nrows));
3069566063dSJacob Faibussowitsch   PetscCall(ISGetLocalSize(iscol, &ncols));
3079566063dSJacob Faibussowitsch   PetscCall(ISGetSize(iscol, &Ncols)); /* global number of columns, size of iscol_local */
308ca3fa75bSLois Curfman McInnes 
309ca3fa75bSLois Curfman McInnes   /* No parallel redistribution currently supported! Should really check each index set
310da81f932SPierre Jolivet      to confirm that it is OK.  ... Currently supports only submatrix same partitioning as
3117eba5e9cSLois Curfman McInnes      original matrix! */
312ca3fa75bSLois Curfman McInnes 
3139566063dSJacob Faibussowitsch   PetscCall(MatGetLocalSize(A, &nlrows, &nlcols));
3149566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(A, &rstart, &rend));
315ca3fa75bSLois Curfman McInnes 
316ca3fa75bSLois Curfman McInnes   /* Check submatrix call */
317ca3fa75bSLois Curfman McInnes   if (scall == MAT_REUSE_MATRIX) {
318e32f2f54SBarry Smith     /* SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Reused submatrix wrong size"); */
3197eba5e9cSLois Curfman McInnes     /* Really need to test rows and column sizes! */
320ca3fa75bSLois Curfman McInnes     newmat = *B;
321ca3fa75bSLois Curfman McInnes   } else {
322ca3fa75bSLois Curfman McInnes     /* Create and fill new matrix */
3239566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &newmat));
3249566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(newmat, nrows, ncols, PETSC_DECIDE, Ncols));
3259566063dSJacob Faibussowitsch     PetscCall(MatSetType(newmat, ((PetscObject)A)->type_name));
3269566063dSJacob Faibussowitsch     PetscCall(MatMPIDenseSetPreallocation(newmat, NULL));
327ca3fa75bSLois Curfman McInnes   }
328ca3fa75bSLois Curfman McInnes 
329ca3fa75bSLois Curfman McInnes   /* Now extract the data pointers and do the copy, column at a time */
330ca3fa75bSLois Curfman McInnes   newmatd = (Mat_MPIDense *)newmat->data;
3319566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(newmatd->A, &bv));
3329566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(mat->A, &v));
3339566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(mat->A, &lda));
3344aa3045dSJed Brown   for (i = 0; i < Ncols; i++) {
335637a0070SStefano Zampini     const PetscScalar *av = v + lda * icol[i];
336ad540459SPierre Jolivet     for (j = 0; j < nrows; j++) *bv++ = av[irow[j] - rstart];
337ca3fa75bSLois Curfman McInnes   }
3389566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(mat->A, &v));
3399566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(newmatd->A, &bv));
340ca3fa75bSLois Curfman McInnes 
341ca3fa75bSLois Curfman McInnes   /* Assemble the matrices so that the correct flags are set */
3429566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(newmat, MAT_FINAL_ASSEMBLY));
3439566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(newmat, MAT_FINAL_ASSEMBLY));
344ca3fa75bSLois Curfman McInnes 
345ca3fa75bSLois Curfman McInnes   /* Free work space */
3469566063dSJacob Faibussowitsch   PetscCall(ISRestoreIndices(isrow, &irow));
3479566063dSJacob Faibussowitsch   PetscCall(ISRestoreIndices(iscol_local, &icol));
3489566063dSJacob Faibussowitsch   PetscCall(ISDestroy(&iscol_local));
349ca3fa75bSLois Curfman McInnes   *B = newmat;
3503ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
351ca3fa75bSLois Curfman McInnes }
352ca3fa75bSLois Curfman McInnes 
35366976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreArray_MPIDense(Mat A, PetscScalar **array)
354d71ae5a4SJacob Faibussowitsch {
35573a71a0fSBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
35673a71a0fSBarry Smith 
3573a40ed3dSBarry Smith   PetscFunctionBegin;
3589566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(a->A, array));
3593ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
360ff14e315SSatish Balay }
361ff14e315SSatish Balay 
362f1e99121SPierre Jolivet static PetscErrorCode MatDenseRestoreArrayRead_MPIDense(Mat A, PetscScalar **array)
363d71ae5a4SJacob Faibussowitsch {
3648572280aSBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
3658572280aSBarry Smith 
3668572280aSBarry Smith   PetscFunctionBegin;
367f1e99121SPierre Jolivet   PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)array));
3683ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3698572280aSBarry Smith }
3708572280aSBarry Smith 
37166976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreArrayWrite_MPIDense(Mat A, PetscScalar **array)
372d71ae5a4SJacob Faibussowitsch {
3736947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
3746947451fSStefano Zampini 
3756947451fSStefano Zampini   PetscFunctionBegin;
3769566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(a->A, array));
3773ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3786947451fSStefano Zampini }
3796947451fSStefano Zampini 
38066976f2fSJacob Faibussowitsch static PetscErrorCode MatAssemblyBegin_MPIDense(Mat mat, MatAssemblyType mode)
381d71ae5a4SJacob Faibussowitsch {
38239ddd567SLois Curfman McInnes   Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
38313f74950SBarry Smith   PetscInt      nstash, reallocs;
3848965ea79SLois Curfman McInnes 
3853a40ed3dSBarry Smith   PetscFunctionBegin;
3863ba16761SJacob Faibussowitsch   if (mdn->donotstash || mat->nooffprocentries) PetscFunctionReturn(PETSC_SUCCESS);
3878965ea79SLois Curfman McInnes 
3889566063dSJacob Faibussowitsch   PetscCall(MatStashScatterBegin_Private(mat, &mat->stash, mat->rmap->range));
3899566063dSJacob Faibussowitsch   PetscCall(MatStashGetInfo_Private(&mat->stash, &nstash, &reallocs));
3909566063dSJacob Faibussowitsch   PetscCall(PetscInfo(mdn->A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
3913ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3928965ea79SLois Curfman McInnes }
3938965ea79SLois Curfman McInnes 
39466976f2fSJacob Faibussowitsch static PetscErrorCode MatAssemblyEnd_MPIDense(Mat mat, MatAssemblyType mode)
395d71ae5a4SJacob Faibussowitsch {
39639ddd567SLois Curfman McInnes   Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
39713f74950SBarry Smith   PetscInt      i, *row, *col, flg, j, rstart, ncols;
39813f74950SBarry Smith   PetscMPIInt   n;
39987828ca2SBarry Smith   PetscScalar  *val;
4008965ea79SLois Curfman McInnes 
4013a40ed3dSBarry Smith   PetscFunctionBegin;
402910cf402Sprj-   if (!mdn->donotstash && !mat->nooffprocentries) {
4038965ea79SLois Curfman McInnes     /*  wait on receives */
4047ef1d9bdSSatish Balay     while (1) {
4059566063dSJacob Faibussowitsch       PetscCall(MatStashScatterGetMesg_Private(&mat->stash, &n, &row, &col, &val, &flg));
4067ef1d9bdSSatish Balay       if (!flg) break;
4078965ea79SLois Curfman McInnes 
4087ef1d9bdSSatish Balay       for (i = 0; i < n;) {
4097ef1d9bdSSatish Balay         /* Now identify the consecutive vals belonging to the same row */
4102205254eSKarl Rupp         for (j = i, rstart = row[j]; j < n; j++) {
4112205254eSKarl Rupp           if (row[j] != rstart) break;
4122205254eSKarl Rupp         }
4137ef1d9bdSSatish Balay         if (j < n) ncols = j - i;
4147ef1d9bdSSatish Balay         else ncols = n - i;
4157ef1d9bdSSatish Balay         /* Now assemble all these values with a single function call */
4169566063dSJacob Faibussowitsch         PetscCall(MatSetValues_MPIDense(mat, 1, row + i, ncols, col + i, val + i, mat->insertmode));
4177ef1d9bdSSatish Balay         i = j;
4188965ea79SLois Curfman McInnes       }
4197ef1d9bdSSatish Balay     }
4209566063dSJacob Faibussowitsch     PetscCall(MatStashScatterEnd_Private(&mat->stash));
421910cf402Sprj-   }
4228965ea79SLois Curfman McInnes 
4239566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(mdn->A, mode));
4249566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(mdn->A, mode));
4253ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4268965ea79SLois Curfman McInnes }
4278965ea79SLois Curfman McInnes 
42866976f2fSJacob Faibussowitsch static PetscErrorCode MatZeroEntries_MPIDense(Mat A)
429d71ae5a4SJacob Faibussowitsch {
43039ddd567SLois Curfman McInnes   Mat_MPIDense *l = (Mat_MPIDense *)A->data;
4313a40ed3dSBarry Smith 
4323a40ed3dSBarry Smith   PetscFunctionBegin;
4339566063dSJacob Faibussowitsch   PetscCall(MatZeroEntries(l->A));
4343ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4358965ea79SLois Curfman McInnes }
4368965ea79SLois Curfman McInnes 
43766976f2fSJacob Faibussowitsch static PetscErrorCode MatZeroRows_MPIDense(Mat A, PetscInt n, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
438d71ae5a4SJacob Faibussowitsch {
43939ddd567SLois Curfman McInnes   Mat_MPIDense *l = (Mat_MPIDense *)A->data;
440637a0070SStefano Zampini   PetscInt      i, len, *lrows;
441637a0070SStefano Zampini 
442637a0070SStefano Zampini   PetscFunctionBegin;
443637a0070SStefano Zampini   /* get locally owned rows */
4449566063dSJacob Faibussowitsch   PetscCall(PetscLayoutMapLocal(A->rmap, n, rows, &len, &lrows, NULL));
445dd8e379bSPierre Jolivet   /* fix right-hand side if needed */
446637a0070SStefano Zampini   if (x && b) {
44797b48c8fSBarry Smith     const PetscScalar *xx;
44897b48c8fSBarry Smith     PetscScalar       *bb;
4498965ea79SLois Curfman McInnes 
4509566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(x, &xx));
4519566063dSJacob Faibussowitsch     PetscCall(VecGetArrayWrite(b, &bb));
452637a0070SStefano Zampini     for (i = 0; i < len; ++i) bb[lrows[i]] = diag * xx[lrows[i]];
4539566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(x, &xx));
4549566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayWrite(b, &bb));
45597b48c8fSBarry Smith   }
4569566063dSJacob Faibussowitsch   PetscCall(MatZeroRows(l->A, len, lrows, 0.0, NULL, NULL));
457e2eb51b1SBarry Smith   if (diag != 0.0) {
458637a0070SStefano Zampini     Vec d;
459b9679d65SBarry Smith 
4609566063dSJacob Faibussowitsch     PetscCall(MatCreateVecs(A, NULL, &d));
4619566063dSJacob Faibussowitsch     PetscCall(VecSet(d, diag));
4629566063dSJacob Faibussowitsch     PetscCall(MatDiagonalSet(A, d, INSERT_VALUES));
4639566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&d));
464b9679d65SBarry Smith   }
4659566063dSJacob Faibussowitsch   PetscCall(PetscFree(lrows));
4663ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4678965ea79SLois Curfman McInnes }
4688965ea79SLois Curfman McInnes 
4690be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMult_SeqDense(Mat, Vec, Vec);
4700be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMultAdd_SeqDense(Mat, Vec, Vec, Vec);
4710be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMultTranspose_SeqDense(Mat, Vec, Vec);
4720be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMultTransposeAdd_SeqDense(Mat, Vec, Vec, Vec);
4730be0d8bdSHansol Suh 
474d016bddeSToby Isaac static PetscErrorCode MatMultColumnRange_MPIDense(Mat mat, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end)
475d016bddeSToby Isaac {
476d016bddeSToby Isaac   Mat_MPIDense      *mdn = (Mat_MPIDense *)mat->data;
477d016bddeSToby Isaac   const PetscScalar *ax;
478d016bddeSToby Isaac   PetscScalar       *ay;
479d016bddeSToby Isaac   PetscMemType       axmtype, aymtype;
480d016bddeSToby Isaac 
481d016bddeSToby Isaac   PetscFunctionBegin;
482d016bddeSToby Isaac   if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
483d016bddeSToby Isaac   PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
484d016bddeSToby Isaac   PetscCall(VecGetArrayWriteAndMemType(mdn->lvec, &ay, &aymtype));
485d016bddeSToby Isaac   PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
486d016bddeSToby Isaac   PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
487d016bddeSToby Isaac   PetscCall(VecRestoreArrayWriteAndMemType(mdn->lvec, &ay));
488d016bddeSToby Isaac   PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
489d016bddeSToby Isaac   PetscUseMethod(mdn->A, "MatMultColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (mdn->A, mdn->lvec, yy, c_start, c_end));
490d016bddeSToby Isaac   PetscFunctionReturn(PETSC_SUCCESS);
491d016bddeSToby Isaac }
492d016bddeSToby Isaac 
49366976f2fSJacob Faibussowitsch static PetscErrorCode MatMult_MPIDense(Mat mat, Vec xx, Vec yy)
494d71ae5a4SJacob Faibussowitsch {
49539ddd567SLois Curfman McInnes   Mat_MPIDense      *mdn = (Mat_MPIDense *)mat->data;
496637a0070SStefano Zampini   const PetscScalar *ax;
497637a0070SStefano Zampini   PetscScalar       *ay;
498d0295fc0SJunchao Zhang   PetscMemType       axmtype, aymtype;
499c456f294SBarry Smith 
5003a40ed3dSBarry Smith   PetscFunctionBegin;
5016f08ad05SPierre Jolivet   if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
5029566063dSJacob Faibussowitsch   PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
503d016bddeSToby Isaac   PetscCall(VecGetArrayWriteAndMemType(mdn->lvec, &ay, &aymtype));
5049566063dSJacob Faibussowitsch   PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
5059566063dSJacob Faibussowitsch   PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
506d016bddeSToby Isaac   PetscCall(VecRestoreArrayWriteAndMemType(mdn->lvec, &ay));
5079566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
5089566063dSJacob Faibussowitsch   PetscCall((*mdn->A->ops->mult)(mdn->A, mdn->lvec, yy));
5093ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
5108965ea79SLois Curfman McInnes }
5118965ea79SLois Curfman McInnes 
5120be0d8bdSHansol Suh static PetscErrorCode MatMultAddColumnRange_MPIDense(Mat mat, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end)
5130be0d8bdSHansol Suh {
5140be0d8bdSHansol Suh   Mat_MPIDense      *mdn = (Mat_MPIDense *)mat->data;
5150be0d8bdSHansol Suh   const PetscScalar *ax;
5160be0d8bdSHansol Suh   PetscScalar       *ay;
5170be0d8bdSHansol Suh   PetscMemType       axmtype, aymtype;
5180be0d8bdSHansol Suh 
5190be0d8bdSHansol Suh   PetscFunctionBegin;
5200be0d8bdSHansol Suh   if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
5210be0d8bdSHansol Suh   PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
5220be0d8bdSHansol Suh   PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype));
5230be0d8bdSHansol Suh   PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
5240be0d8bdSHansol Suh   PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
5250be0d8bdSHansol Suh   PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay));
5260be0d8bdSHansol Suh   PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
5270be0d8bdSHansol Suh   PetscUseMethod(mdn->A, "MatMultAddColumnRange_C", (Mat, Vec, Vec, Vec, PetscInt, PetscInt), (mdn->A, mdn->lvec, yy, zz, c_start, c_end));
5280be0d8bdSHansol Suh   PetscFunctionReturn(PETSC_SUCCESS);
5290be0d8bdSHansol Suh }
5300be0d8bdSHansol Suh 
53166976f2fSJacob Faibussowitsch static PetscErrorCode MatMultAdd_MPIDense(Mat mat, Vec xx, Vec yy, Vec zz)
532d71ae5a4SJacob Faibussowitsch {
53339ddd567SLois Curfman McInnes   Mat_MPIDense      *mdn = (Mat_MPIDense *)mat->data;
534637a0070SStefano Zampini   const PetscScalar *ax;
535637a0070SStefano Zampini   PetscScalar       *ay;
536d0295fc0SJunchao Zhang   PetscMemType       axmtype, aymtype;
537c456f294SBarry Smith 
5383a40ed3dSBarry Smith   PetscFunctionBegin;
5396f08ad05SPierre Jolivet   if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
5409566063dSJacob Faibussowitsch   PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
5419566063dSJacob Faibussowitsch   PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype));
5429566063dSJacob Faibussowitsch   PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
5439566063dSJacob Faibussowitsch   PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
5449566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay));
5459566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
5469566063dSJacob Faibussowitsch   PetscCall((*mdn->A->ops->multadd)(mdn->A, mdn->lvec, yy, zz));
5473ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
5488965ea79SLois Curfman McInnes }
5498965ea79SLois Curfman McInnes 
5500be0d8bdSHansol Suh static PetscErrorCode MatMultHermitianTransposeColumnRange_MPIDense(Mat A, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end)
5510be0d8bdSHansol Suh {
5520be0d8bdSHansol Suh   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
5530be0d8bdSHansol Suh   const PetscScalar *ax;
5540be0d8bdSHansol Suh   PetscScalar       *ay;
5550be0d8bdSHansol Suh   PetscMemType       axmtype, aymtype;
556d016bddeSToby Isaac   PetscInt           r_start, r_end;
557d016bddeSToby Isaac   PetscInt           c_start_local, c_end_local;
5580be0d8bdSHansol Suh 
5590be0d8bdSHansol Suh   PetscFunctionBegin;
5600be0d8bdSHansol Suh   if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
561d016bddeSToby Isaac   PetscCall(VecZeroEntries(a->lvec));
562d016bddeSToby Isaac   PetscCall(VecGetOwnershipRange(yy, &r_start, &r_end));
563d016bddeSToby Isaac   c_start_local = PetscMax(c_start, r_start);
564d016bddeSToby Isaac   c_end_local   = PetscMin(c_end, r_end);
565d016bddeSToby Isaac   PetscCall(VecGetArrayAndMemType(yy, &ay, &aymtype));
566d016bddeSToby Isaac   if (c_end_local > c_start_local) {
567d016bddeSToby Isaac     if (PetscMemTypeHost(aymtype)) {
568d016bddeSToby Isaac       PetscCall(PetscArrayzero(&ay[c_start_local], (size_t)(c_end_local - c_start_local)));
569d016bddeSToby Isaac     } else {
570d016bddeSToby Isaac       PetscCall(PetscDeviceRegisterMemory(ay, aymtype, sizeof(*ay) * ((size_t)(r_end - r_start))));
571d016bddeSToby Isaac       PetscCall(PetscDeviceArrayZero(NULL, &ay[c_start_local], (size_t)(c_end_local - c_start_local)));
572d016bddeSToby Isaac     }
573d016bddeSToby Isaac   }
5740be0d8bdSHansol Suh   PetscUseMethod(a->A, "MatMultHermitianTransposeColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (a->A, xx, a->lvec, c_start, c_end));
5750be0d8bdSHansol Suh   PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
5760be0d8bdSHansol Suh   PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
5770be0d8bdSHansol Suh   PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
5780be0d8bdSHansol Suh   PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
5790be0d8bdSHansol Suh   PetscCall(VecRestoreArrayAndMemType(yy, &ay));
5800be0d8bdSHansol Suh   PetscFunctionReturn(PETSC_SUCCESS);
5810be0d8bdSHansol Suh }
5820be0d8bdSHansol Suh 
583459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTransposeKernel_MPIDense(Mat A, Vec xx, Vec yy, PetscBool herm)
584d71ae5a4SJacob Faibussowitsch {
585096963f5SLois Curfman McInnes   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
586637a0070SStefano Zampini   const PetscScalar *ax;
587637a0070SStefano Zampini   PetscScalar       *ay;
588d0295fc0SJunchao Zhang   PetscMemType       axmtype, aymtype;
5894d86920dSPierre Jolivet 
5903a40ed3dSBarry Smith   PetscFunctionBegin;
5916f08ad05SPierre Jolivet   if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
5929566063dSJacob Faibussowitsch   PetscCall(VecSet(yy, 0.0));
593459e8d23SBlanca Mellado Pinto   if (herm) PetscCall((*a->A->ops->multhermitiantranspose)(a->A, xx, a->lvec));
594459e8d23SBlanca Mellado Pinto   else PetscCall((*a->A->ops->multtranspose)(a->A, xx, a->lvec));
5959566063dSJacob Faibussowitsch   PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
5969566063dSJacob Faibussowitsch   PetscCall(VecGetArrayAndMemType(yy, &ay, &aymtype));
5979566063dSJacob Faibussowitsch   PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
5989566063dSJacob Faibussowitsch   PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
5999566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
6009566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayAndMemType(yy, &ay));
6013ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
602096963f5SLois Curfman McInnes }
603096963f5SLois Curfman McInnes 
6040be0d8bdSHansol Suh static PetscErrorCode MatMultHermitianTransposeAddColumnRange_MPIDense(Mat A, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end)
6050be0d8bdSHansol Suh {
6060be0d8bdSHansol Suh   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
6070be0d8bdSHansol Suh   const PetscScalar *ax;
6080be0d8bdSHansol Suh   PetscScalar       *ay;
6090be0d8bdSHansol Suh   PetscMemType       axmtype, aymtype;
6100be0d8bdSHansol Suh 
6110be0d8bdSHansol Suh   PetscFunctionBegin;
6120be0d8bdSHansol Suh   if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
6130be0d8bdSHansol Suh   PetscCall(VecCopy(yy, zz));
614d016bddeSToby Isaac   PetscCall(VecZeroEntries(a->lvec));
6150be0d8bdSHansol Suh   PetscUseMethod(a->A, "MatMultHermitianTransposeColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (a->A, xx, a->lvec, c_start, c_end));
6160be0d8bdSHansol Suh   PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
6170be0d8bdSHansol Suh   PetscCall(VecGetArrayAndMemType(zz, &ay, &aymtype));
6180be0d8bdSHansol Suh   PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
6190be0d8bdSHansol Suh   PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
6200be0d8bdSHansol Suh   PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
6210be0d8bdSHansol Suh   PetscCall(VecRestoreArrayAndMemType(zz, &ay));
6220be0d8bdSHansol Suh   PetscFunctionReturn(PETSC_SUCCESS);
6230be0d8bdSHansol Suh }
6240be0d8bdSHansol Suh 
625459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTransposeAddKernel_MPIDense(Mat A, Vec xx, Vec yy, Vec zz, PetscBool herm)
626d71ae5a4SJacob Faibussowitsch {
627096963f5SLois Curfman McInnes   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
628637a0070SStefano Zampini   const PetscScalar *ax;
629637a0070SStefano Zampini   PetscScalar       *ay;
630d0295fc0SJunchao Zhang   PetscMemType       axmtype, aymtype;
631096963f5SLois Curfman McInnes 
6323a40ed3dSBarry Smith   PetscFunctionBegin;
6336f08ad05SPierre Jolivet   if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
6349566063dSJacob Faibussowitsch   PetscCall(VecCopy(yy, zz));
635459e8d23SBlanca Mellado Pinto   if (herm) PetscCall((*a->A->ops->multhermitiantranspose)(a->A, xx, a->lvec));
636459e8d23SBlanca Mellado Pinto   else PetscCall((*a->A->ops->multtranspose)(a->A, xx, a->lvec));
6379566063dSJacob Faibussowitsch   PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
6389566063dSJacob Faibussowitsch   PetscCall(VecGetArrayAndMemType(zz, &ay, &aymtype));
6399566063dSJacob Faibussowitsch   PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
6409566063dSJacob Faibussowitsch   PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
6419566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
6429566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayAndMemType(zz, &ay));
6433ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
644096963f5SLois Curfman McInnes }
645096963f5SLois Curfman McInnes 
646459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTranspose_MPIDense(Mat A, Vec xx, Vec yy)
647459e8d23SBlanca Mellado Pinto {
648459e8d23SBlanca Mellado Pinto   PetscFunctionBegin;
649459e8d23SBlanca Mellado Pinto   PetscCall(MatMultTransposeKernel_MPIDense(A, xx, yy, PETSC_FALSE));
650459e8d23SBlanca Mellado Pinto   PetscFunctionReturn(PETSC_SUCCESS);
651459e8d23SBlanca Mellado Pinto }
652459e8d23SBlanca Mellado Pinto 
653459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTransposeAdd_MPIDense(Mat A, Vec xx, Vec yy, Vec zz)
654459e8d23SBlanca Mellado Pinto {
655459e8d23SBlanca Mellado Pinto   PetscFunctionBegin;
656459e8d23SBlanca Mellado Pinto   PetscCall(MatMultTransposeAddKernel_MPIDense(A, xx, yy, zz, PETSC_FALSE));
657459e8d23SBlanca Mellado Pinto   PetscFunctionReturn(PETSC_SUCCESS);
658459e8d23SBlanca Mellado Pinto }
659459e8d23SBlanca Mellado Pinto 
660459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultHermitianTranspose_MPIDense(Mat A, Vec xx, Vec yy)
661459e8d23SBlanca Mellado Pinto {
662459e8d23SBlanca Mellado Pinto   PetscFunctionBegin;
663459e8d23SBlanca Mellado Pinto   PetscCall(MatMultTransposeKernel_MPIDense(A, xx, yy, PETSC_TRUE));
664459e8d23SBlanca Mellado Pinto   PetscFunctionReturn(PETSC_SUCCESS);
665459e8d23SBlanca Mellado Pinto }
666459e8d23SBlanca Mellado Pinto 
667459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultHermitianTransposeAdd_MPIDense(Mat A, Vec xx, Vec yy, Vec zz)
668459e8d23SBlanca Mellado Pinto {
669459e8d23SBlanca Mellado Pinto   PetscFunctionBegin;
670459e8d23SBlanca Mellado Pinto   PetscCall(MatMultTransposeAddKernel_MPIDense(A, xx, yy, zz, PETSC_TRUE));
671459e8d23SBlanca Mellado Pinto   PetscFunctionReturn(PETSC_SUCCESS);
672459e8d23SBlanca Mellado Pinto }
673459e8d23SBlanca Mellado Pinto 
674d71ae5a4SJacob Faibussowitsch PetscErrorCode MatGetDiagonal_MPIDense(Mat A, Vec v)
675d71ae5a4SJacob Faibussowitsch {
67639ddd567SLois Curfman McInnes   Mat_MPIDense      *a = (Mat_MPIDense *)A->data;
67779560b15SJacob Faibussowitsch   PetscInt           lda, len, i, nl, ng, m = A->rmap->n, radd;
67879560b15SJacob Faibussowitsch   PetscScalar       *x;
679637a0070SStefano Zampini   const PetscScalar *av;
680ed3cc1f0SBarry Smith 
6813a40ed3dSBarry Smith   PetscFunctionBegin;
6829566063dSJacob Faibussowitsch   PetscCall(VecGetArray(v, &x));
68379560b15SJacob Faibussowitsch   PetscCall(VecGetSize(v, &ng));
68479560b15SJacob Faibussowitsch   PetscCheck(ng == A->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Nonconforming mat and vec");
68579560b15SJacob Faibussowitsch   PetscCall(VecGetLocalSize(v, &nl));
686d0f46423SBarry Smith   len  = PetscMin(a->A->rmap->n, a->A->cmap->n);
687d0f46423SBarry Smith   radd = A->rmap->rstart * m;
6889566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, &av));
6899566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
690ad540459SPierre Jolivet   for (i = 0; i < len; i++) x[i] = av[radd + i * lda + i];
6919566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, &av));
6920be0d8bdSHansol Suh   if (nl - i > 0) PetscCall(PetscArrayzero(x + i, nl - i));
6939566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(v, &x));
6943ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
6958965ea79SLois Curfman McInnes }
6968965ea79SLois Curfman McInnes 
69766976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MPIDense(Mat mat)
698d71ae5a4SJacob Faibussowitsch {
6993501a2bdSLois Curfman McInnes   Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
700ed3cc1f0SBarry Smith 
7013a40ed3dSBarry Smith   PetscFunctionBegin;
7023ba16761SJacob Faibussowitsch   PetscCall(PetscLogObjectState((PetscObject)mat, "Rows=%" PetscInt_FMT ", Cols=%" PetscInt_FMT, mat->rmap->N, mat->cmap->N));
7039566063dSJacob Faibussowitsch   PetscCall(MatStashDestroy_Private(&mat->stash));
70428b400f6SJacob Faibussowitsch   PetscCheck(!mdn->vecinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
70528b400f6SJacob Faibussowitsch   PetscCheck(!mdn->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
7069566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&mdn->A));
7079566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&mdn->lvec));
7089566063dSJacob Faibussowitsch   PetscCall(PetscSFDestroy(&mdn->Mvctx));
7099566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&mdn->cvec));
7109566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&mdn->cmat));
71101b82886SBarry Smith 
7129566063dSJacob Faibussowitsch   PetscCall(PetscFree(mat->data));
7139566063dSJacob Faibussowitsch   PetscCall(PetscObjectChangeTypeName((PetscObject)mat, NULL));
7148baccfbdSHong Zhang 
7159566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", NULL));
7169566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", NULL));
7179566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", NULL));
7189566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", NULL));
7199566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", NULL));
7209566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", NULL));
7219566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", NULL));
7229566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", NULL));
7239566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", NULL));
7249566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", NULL));
7259566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", NULL));
7269566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", NULL));
7279566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", NULL));
7288baccfbdSHong Zhang #if defined(PETSC_HAVE_ELEMENTAL)
7299566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", NULL));
7308baccfbdSHong Zhang #endif
731d24d4204SJose E. Roman #if defined(PETSC_HAVE_SCALAPACK)
7329566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", NULL));
733d24d4204SJose E. Roman #endif
7349566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", NULL));
7359566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", NULL));
7369566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", NULL));
7376718818eSStefano Zampini #if defined(PETSC_HAVE_CUDA)
7389566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", NULL));
7399566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", NULL));
7402e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", NULL));
7412e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidensecuda_mpidense_C", NULL));
7422e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidensecuda_C", NULL));
7432e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidensecuda_C", NULL));
7442e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensecuda_mpiaij_C", NULL));
7452e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensecuda_mpiaijcusparse_C", NULL));
7462e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArray_C", NULL));
7472e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArrayRead_C", NULL));
7482e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArrayWrite_C", NULL));
7492e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArray_C", NULL));
7502e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArrayRead_C", NULL));
7512e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArrayWrite_C", NULL));
7522e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAPlaceArray_C", NULL));
7532e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAResetArray_C", NULL));
7542e956fe4SStefano Zampini   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAReplaceArray_C", NULL));
7553d9668e3SJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDASetPreallocation_C", NULL));
7566718818eSStefano Zampini #endif
75747d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
75847d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", NULL));
75947d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", NULL));
76047d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", NULL));
76147d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidensehip_mpidense_C", NULL));
76247d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidensehip_C", NULL));
76347d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidensehip_C", NULL));
76447d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensehip_mpiaij_C", NULL));
76547d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensehip_mpiaijhipsparse_C", NULL));
76647d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArray_C", NULL));
76747d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArrayRead_C", NULL));
76847d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArrayWrite_C", NULL));
76947d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArray_C", NULL));
77047d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArrayRead_C", NULL));
77147d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArrayWrite_C", NULL));
77247d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPPlaceArray_C", NULL));
77347d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPResetArray_C", NULL));
77447d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPReplaceArray_C", NULL));
7753d9668e3SJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPSetPreallocation_C", NULL));
77647d993e7Ssuyashtn #endif
7779566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", NULL));
7789566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", NULL));
7799566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", NULL));
7809566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", NULL));
7819566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", NULL));
7829566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", NULL));
7839566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", NULL));
7849566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", NULL));
7859566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", NULL));
7869566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", NULL));
787d016bddeSToby Isaac   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultColumnRange_C", NULL));
7880be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", NULL));
7890be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", NULL));
7900be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", NULL));
791f8103e6bSStefano Zampini 
792f8103e6bSStefano Zampini   PetscCall(PetscObjectCompose((PetscObject)mat, "DiagonalBlock", NULL));
7933ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
7948965ea79SLois Curfman McInnes }
79539ddd567SLois Curfman McInnes 
7969804daf3SBarry Smith #include <petscdraw.h>
797d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatView_MPIDense_ASCIIorDraworSocket(Mat mat, PetscViewer viewer)
798d71ae5a4SJacob Faibussowitsch {
79939ddd567SLois Curfman McInnes   Mat_MPIDense     *mdn = (Mat_MPIDense *)mat->data;
800616b8fbbSStefano Zampini   PetscMPIInt       rank;
80119fd82e9SBarry Smith   PetscViewerType   vtype;
802ace3abfcSBarry Smith   PetscBool         iascii, isdraw;
803b0a32e0cSBarry Smith   PetscViewer       sviewer;
804f3ef73ceSBarry Smith   PetscViewerFormat format;
8058965ea79SLois Curfman McInnes 
8063a40ed3dSBarry Smith   PetscFunctionBegin;
8079566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mat), &rank));
8089566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
8099566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
81032077d6dSBarry Smith   if (iascii) {
8119566063dSJacob Faibussowitsch     PetscCall(PetscViewerGetType(viewer, &vtype));
8129566063dSJacob Faibussowitsch     PetscCall(PetscViewerGetFormat(viewer, &format));
813456192e2SBarry Smith     if (format == PETSC_VIEWER_ASCII_INFO_DETAIL) {
8144e220ebcSLois Curfman McInnes       MatInfo info;
8159566063dSJacob Faibussowitsch       PetscCall(MatGetInfo(mat, MAT_LOCAL, &info));
8169566063dSJacob Faibussowitsch       PetscCall(PetscViewerASCIIPushSynchronized(viewer));
8179371c9d4SSatish 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,
8189371c9d4SSatish Balay                                                    (PetscInt)info.memory));
8199566063dSJacob Faibussowitsch       PetscCall(PetscViewerFlush(viewer));
8209566063dSJacob Faibussowitsch       PetscCall(PetscViewerASCIIPopSynchronized(viewer));
8216f08ad05SPierre Jolivet       if (mdn->Mvctx) PetscCall(PetscSFView(mdn->Mvctx, viewer));
8223ba16761SJacob Faibussowitsch       PetscFunctionReturn(PETSC_SUCCESS);
823fb9695e5SSatish Balay     } else if (format == PETSC_VIEWER_ASCII_INFO) {
8243ba16761SJacob Faibussowitsch       PetscFunctionReturn(PETSC_SUCCESS);
8258965ea79SLois Curfman McInnes     }
826f1af5d2fSBarry Smith   } else if (isdraw) {
827b0a32e0cSBarry Smith     PetscDraw draw;
828ace3abfcSBarry Smith     PetscBool isnull;
829f1af5d2fSBarry Smith 
8309566063dSJacob Faibussowitsch     PetscCall(PetscViewerDrawGetDraw(viewer, 0, &draw));
8319566063dSJacob Faibussowitsch     PetscCall(PetscDrawIsNull(draw, &isnull));
8323ba16761SJacob Faibussowitsch     if (isnull) PetscFunctionReturn(PETSC_SUCCESS);
833f1af5d2fSBarry Smith   }
83477ed5343SBarry Smith 
8357da1fb6eSBarry Smith   {
8368965ea79SLois Curfman McInnes     /* assemble the entire matrix onto first processor. */
8378965ea79SLois Curfman McInnes     Mat          A;
838d0f46423SBarry Smith     PetscInt     M = mat->rmap->N, N = mat->cmap->N, m, row, i, nz;
839ba8c8a56SBarry Smith     PetscInt    *cols;
840ba8c8a56SBarry Smith     PetscScalar *vals;
8418965ea79SLois Curfman McInnes 
8429566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)mat), &A));
843dd400576SPatrick Sanan     if (rank == 0) {
8449566063dSJacob Faibussowitsch       PetscCall(MatSetSizes(A, M, N, M, N));
8453a40ed3dSBarry Smith     } else {
8469566063dSJacob Faibussowitsch       PetscCall(MatSetSizes(A, 0, 0, M, N));
8478965ea79SLois Curfman McInnes     }
8487adad957SLisandro Dalcin     /* Since this is a temporary matrix, MATMPIDENSE instead of ((PetscObject)A)->type_name here is probably acceptable. */
8499566063dSJacob Faibussowitsch     PetscCall(MatSetType(A, MATMPIDENSE));
8509566063dSJacob Faibussowitsch     PetscCall(MatMPIDenseSetPreallocation(A, NULL));
8518965ea79SLois Curfman McInnes 
85239ddd567SLois Curfman McInnes     /* Copy the matrix ... This isn't the most efficient means,
85339ddd567SLois Curfman McInnes        but it's quick for now */
85451022da4SBarry Smith     A->insertmode = INSERT_VALUES;
8552205254eSKarl Rupp 
8562205254eSKarl Rupp     row = mat->rmap->rstart;
8572205254eSKarl Rupp     m   = mdn->A->rmap->n;
8588965ea79SLois Curfman McInnes     for (i = 0; i < m; i++) {
8599566063dSJacob Faibussowitsch       PetscCall(MatGetRow_MPIDense(mat, row, &nz, &cols, &vals));
8609566063dSJacob Faibussowitsch       PetscCall(MatSetValues_MPIDense(A, 1, &row, nz, cols, vals, INSERT_VALUES));
8619566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow_MPIDense(mat, row, &nz, &cols, &vals));
86239ddd567SLois Curfman McInnes       row++;
8638965ea79SLois Curfman McInnes     }
8648965ea79SLois Curfman McInnes 
8659566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
8669566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
8679566063dSJacob Faibussowitsch     PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
868dd400576SPatrick Sanan     if (rank == 0) {
869f4f49eeaSPierre Jolivet       PetscCall(PetscObjectSetName((PetscObject)((Mat_MPIDense *)A->data)->A, ((PetscObject)mat)->name));
870f4f49eeaSPierre Jolivet       PetscCall(MatView_SeqDense(((Mat_MPIDense *)A->data)->A, sviewer));
8718965ea79SLois Curfman McInnes     }
8729566063dSJacob Faibussowitsch     PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
8739566063dSJacob Faibussowitsch     PetscCall(MatDestroy(&A));
8748965ea79SLois Curfman McInnes   }
8753ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
8768965ea79SLois Curfman McInnes }
8778965ea79SLois Curfman McInnes 
87866976f2fSJacob Faibussowitsch static PetscErrorCode MatView_MPIDense(Mat mat, PetscViewer viewer)
879d71ae5a4SJacob Faibussowitsch {
880ace3abfcSBarry Smith   PetscBool iascii, isbinary, isdraw, issocket;
8818965ea79SLois Curfman McInnes 
882433994e6SBarry Smith   PetscFunctionBegin;
8839566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
8849566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
8859566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERSOCKET, &issocket));
8869566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
8870f5bd95cSBarry Smith 
88832077d6dSBarry Smith   if (iascii || issocket || isdraw) {
8899566063dSJacob Faibussowitsch     PetscCall(MatView_MPIDense_ASCIIorDraworSocket(mat, viewer));
8901baa6e33SBarry Smith   } else if (isbinary) PetscCall(MatView_Dense_Binary(mat, viewer));
8913ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
8928965ea79SLois Curfman McInnes }
8938965ea79SLois Curfman McInnes 
89466976f2fSJacob Faibussowitsch static PetscErrorCode MatGetInfo_MPIDense(Mat A, MatInfoType flag, MatInfo *info)
895d71ae5a4SJacob Faibussowitsch {
8963501a2bdSLois Curfman McInnes   Mat_MPIDense  *mat = (Mat_MPIDense *)A->data;
8973501a2bdSLois Curfman McInnes   Mat            mdn = mat->A;
8983966268fSBarry Smith   PetscLogDouble isend[5], irecv[5];
8998965ea79SLois Curfman McInnes 
9003a40ed3dSBarry Smith   PetscFunctionBegin;
9014e220ebcSLois Curfman McInnes   info->block_size = 1.0;
9022205254eSKarl Rupp 
9039566063dSJacob Faibussowitsch   PetscCall(MatGetInfo(mdn, MAT_LOCAL, info));
9042205254eSKarl Rupp 
9059371c9d4SSatish Balay   isend[0] = info->nz_used;
9069371c9d4SSatish Balay   isend[1] = info->nz_allocated;
9079371c9d4SSatish Balay   isend[2] = info->nz_unneeded;
9089371c9d4SSatish Balay   isend[3] = info->memory;
9099371c9d4SSatish Balay   isend[4] = info->mallocs;
9108965ea79SLois Curfman McInnes   if (flag == MAT_LOCAL) {
9114e220ebcSLois Curfman McInnes     info->nz_used      = isend[0];
9124e220ebcSLois Curfman McInnes     info->nz_allocated = isend[1];
9134e220ebcSLois Curfman McInnes     info->nz_unneeded  = isend[2];
9144e220ebcSLois Curfman McInnes     info->memory       = isend[3];
9154e220ebcSLois Curfman McInnes     info->mallocs      = isend[4];
9168965ea79SLois Curfman McInnes   } else if (flag == MAT_GLOBAL_MAX) {
917462c564dSBarry Smith     PetscCallMPI(MPIU_Allreduce(isend, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_MAX, PetscObjectComm((PetscObject)A)));
9182205254eSKarl Rupp 
9194e220ebcSLois Curfman McInnes     info->nz_used      = irecv[0];
9204e220ebcSLois Curfman McInnes     info->nz_allocated = irecv[1];
9214e220ebcSLois Curfman McInnes     info->nz_unneeded  = irecv[2];
9224e220ebcSLois Curfman McInnes     info->memory       = irecv[3];
9234e220ebcSLois Curfman McInnes     info->mallocs      = irecv[4];
9248965ea79SLois Curfman McInnes   } else if (flag == MAT_GLOBAL_SUM) {
925462c564dSBarry Smith     PetscCallMPI(MPIU_Allreduce(isend, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_SUM, PetscObjectComm((PetscObject)A)));
9262205254eSKarl Rupp 
9274e220ebcSLois Curfman McInnes     info->nz_used      = irecv[0];
9284e220ebcSLois Curfman McInnes     info->nz_allocated = irecv[1];
9294e220ebcSLois Curfman McInnes     info->nz_unneeded  = irecv[2];
9304e220ebcSLois Curfman McInnes     info->memory       = irecv[3];
9314e220ebcSLois Curfman McInnes     info->mallocs      = irecv[4];
9328965ea79SLois Curfman McInnes   }
9334e220ebcSLois Curfman McInnes   info->fill_ratio_given  = 0; /* no parallel LU/ILU/Cholesky */
9344e220ebcSLois Curfman McInnes   info->fill_ratio_needed = 0;
9354e220ebcSLois Curfman McInnes   info->factor_mallocs    = 0;
9363ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9378965ea79SLois Curfman McInnes }
9388965ea79SLois Curfman McInnes 
93966976f2fSJacob Faibussowitsch static PetscErrorCode MatSetOption_MPIDense(Mat A, MatOption op, PetscBool flg)
940d71ae5a4SJacob Faibussowitsch {
94139ddd567SLois Curfman McInnes   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
9428965ea79SLois Curfman McInnes 
9433a40ed3dSBarry Smith   PetscFunctionBegin;
94412c028f9SKris Buschelman   switch (op) {
945512a5fc5SBarry Smith   case MAT_NEW_NONZERO_LOCATIONS:
94612c028f9SKris Buschelman   case MAT_NEW_NONZERO_LOCATION_ERR:
94712c028f9SKris Buschelman   case MAT_NEW_NONZERO_ALLOCATION_ERR:
94843674050SBarry Smith     MatCheckPreallocated(A, 1);
9499566063dSJacob Faibussowitsch     PetscCall(MatSetOption(a->A, op, flg));
95012c028f9SKris Buschelman     break;
95112c028f9SKris Buschelman   case MAT_ROW_ORIENTED:
95243674050SBarry Smith     MatCheckPreallocated(A, 1);
9534e0d8c25SBarry Smith     a->roworiented = flg;
9549566063dSJacob Faibussowitsch     PetscCall(MatSetOption(a->A, op, flg));
95512c028f9SKris Buschelman     break;
956d71ae5a4SJacob Faibussowitsch   case MAT_IGNORE_OFF_PROC_ENTRIES:
957d71ae5a4SJacob Faibussowitsch     a->donotstash = flg;
958d71ae5a4SJacob Faibussowitsch     break;
95977e54ba9SKris Buschelman   case MAT_SYMMETRIC:
96077e54ba9SKris Buschelman   case MAT_STRUCTURALLY_SYMMETRIC:
9619a4540c5SBarry Smith   case MAT_HERMITIAN:
9629a4540c5SBarry Smith   case MAT_SYMMETRY_ETERNAL:
963b94d7dedSBarry Smith   case MAT_STRUCTURAL_SYMMETRY_ETERNAL:
964b94d7dedSBarry Smith   case MAT_SPD:
965b94d7dedSBarry Smith   case MAT_SPD_ETERNAL:
966b94d7dedSBarry Smith     /* if the diagonal matrix is square it inherits some of the properties above */
9679725ddffSPierre Jolivet     if (a->A && A->rmap->n == A->cmap->n) PetscCall(MatSetOption(a->A, op, flg));
96877e54ba9SKris Buschelman     break;
969d71ae5a4SJacob Faibussowitsch   default:
970888c827cSStefano Zampini     break;
9713a40ed3dSBarry Smith   }
9723ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
9738965ea79SLois Curfman McInnes }
9748965ea79SLois Curfman McInnes 
97566976f2fSJacob Faibussowitsch static PetscErrorCode MatDiagonalScale_MPIDense(Mat A, Vec ll, Vec rr)
976d71ae5a4SJacob Faibussowitsch {
9775b2fa520SLois Curfman McInnes   Mat_MPIDense      *mdn = (Mat_MPIDense *)A->data;
978637a0070SStefano Zampini   const PetscScalar *l;
979637a0070SStefano Zampini   PetscScalar        x, *v, *vv, *r;
980637a0070SStefano Zampini   PetscInt           i, j, s2a, s3a, s2, s3, m = mdn->A->rmap->n, n = mdn->A->cmap->n, lda;
9815b2fa520SLois Curfman McInnes 
9825b2fa520SLois Curfman McInnes   PetscFunctionBegin;
9839566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(mdn->A, &vv));
9849566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(mdn->A, &lda));
9859566063dSJacob Faibussowitsch   PetscCall(MatGetLocalSize(A, &s2, &s3));
9865b2fa520SLois Curfman McInnes   if (ll) {
9879566063dSJacob Faibussowitsch     PetscCall(VecGetLocalSize(ll, &s2a));
98808401ef6SPierre Jolivet     PetscCheck(s2a == s2, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Left scaling vector non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT, s2a, s2);
9899566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(ll, &l));
9905b2fa520SLois Curfman McInnes     for (i = 0; i < m; i++) {
9915b2fa520SLois Curfman McInnes       x = l[i];
992637a0070SStefano Zampini       v = vv + i;
9939371c9d4SSatish Balay       for (j = 0; j < n; j++) {
9949371c9d4SSatish Balay         (*v) *= x;
9959371c9d4SSatish Balay         v += lda;
9969371c9d4SSatish Balay       }
9975b2fa520SLois Curfman McInnes     }
9989566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(ll, &l));
9999566063dSJacob Faibussowitsch     PetscCall(PetscLogFlops(1.0 * n * m));
10005b2fa520SLois Curfman McInnes   }
10015b2fa520SLois Curfman McInnes   if (rr) {
1002637a0070SStefano Zampini     const PetscScalar *ar;
1003637a0070SStefano Zampini 
10049566063dSJacob Faibussowitsch     PetscCall(VecGetLocalSize(rr, &s3a));
100508401ef6SPierre Jolivet     PetscCheck(s3a == s3, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Right scaling vec non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT ".", s3a, s3);
10069566063dSJacob Faibussowitsch     PetscCall(VecGetArrayRead(rr, &ar));
10076f08ad05SPierre Jolivet     if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
10089566063dSJacob Faibussowitsch     PetscCall(VecGetArray(mdn->lvec, &r));
10099566063dSJacob Faibussowitsch     PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, ar, r, MPI_REPLACE));
10109566063dSJacob Faibussowitsch     PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ar, r, MPI_REPLACE));
10119566063dSJacob Faibussowitsch     PetscCall(VecRestoreArrayRead(rr, &ar));
10125b2fa520SLois Curfman McInnes     for (i = 0; i < n; i++) {
10135b2fa520SLois Curfman McInnes       x = r[i];
1014637a0070SStefano Zampini       v = vv + i * lda;
10152205254eSKarl Rupp       for (j = 0; j < m; j++) (*v++) *= x;
10165b2fa520SLois Curfman McInnes     }
10179566063dSJacob Faibussowitsch     PetscCall(VecRestoreArray(mdn->lvec, &r));
10189566063dSJacob Faibussowitsch     PetscCall(PetscLogFlops(1.0 * n * m));
10195b2fa520SLois Curfman McInnes   }
10209566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(mdn->A, &vv));
10213ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
10225b2fa520SLois Curfman McInnes }
10235b2fa520SLois Curfman McInnes 
102466976f2fSJacob Faibussowitsch static PetscErrorCode MatNorm_MPIDense(Mat A, NormType type, PetscReal *nrm)
1025d71ae5a4SJacob Faibussowitsch {
10263501a2bdSLois Curfman McInnes   Mat_MPIDense      *mdn = (Mat_MPIDense *)A->data;
102713f74950SBarry Smith   PetscInt           i, j;
1028e91c04dfSPierre Jolivet   PetscMPIInt        size;
1029329f5518SBarry Smith   PetscReal          sum = 0.0;
1030637a0070SStefano Zampini   const PetscScalar *av, *v;
10313501a2bdSLois Curfman McInnes 
10323a40ed3dSBarry Smith   PetscFunctionBegin;
10339566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(mdn->A, &av));
1034637a0070SStefano Zampini   v = av;
10359566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
1036616b8fbbSStefano Zampini   if (size == 1) {
10379566063dSJacob Faibussowitsch     PetscCall(MatNorm(mdn->A, type, nrm));
10383501a2bdSLois Curfman McInnes   } else {
10393501a2bdSLois Curfman McInnes     if (type == NORM_FROBENIUS) {
1040d0f46423SBarry Smith       for (i = 0; i < mdn->A->cmap->n * mdn->A->rmap->n; i++) {
10419371c9d4SSatish Balay         sum += PetscRealPart(PetscConj(*v) * (*v));
10429371c9d4SSatish Balay         v++;
10433501a2bdSLois Curfman McInnes       }
1044462c564dSBarry Smith       PetscCallMPI(MPIU_Allreduce(&sum, nrm, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
10458f1a2a5eSBarry Smith       *nrm = PetscSqrtReal(*nrm);
10469566063dSJacob Faibussowitsch       PetscCall(PetscLogFlops(2.0 * mdn->A->cmap->n * mdn->A->rmap->n));
10473a40ed3dSBarry Smith     } else if (type == NORM_1) {
1048e91c04dfSPierre Jolivet       PetscReal *tmp;
1049e91c04dfSPierre Jolivet 
1050e91c04dfSPierre Jolivet       PetscCall(PetscCalloc1(A->cmap->N, &tmp));
1051064f8208SBarry Smith       *nrm = 0.0;
1052637a0070SStefano Zampini       v    = av;
1053d0f46423SBarry Smith       for (j = 0; j < mdn->A->cmap->n; j++) {
1054d0f46423SBarry Smith         for (i = 0; i < mdn->A->rmap->n; i++) {
10559371c9d4SSatish Balay           tmp[j] += PetscAbsScalar(*v);
10569371c9d4SSatish Balay           v++;
10573501a2bdSLois Curfman McInnes         }
10583501a2bdSLois Curfman McInnes       }
1059e91c04dfSPierre Jolivet       PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, tmp, A->cmap->N, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
1060d0f46423SBarry Smith       for (j = 0; j < A->cmap->N; j++) {
1061e91c04dfSPierre Jolivet         if (tmp[j] > *nrm) *nrm = tmp[j];
10623501a2bdSLois Curfman McInnes       }
1063e91c04dfSPierre Jolivet       PetscCall(PetscFree(tmp));
10649566063dSJacob Faibussowitsch       PetscCall(PetscLogFlops(A->cmap->n * A->rmap->n));
10653a40ed3dSBarry Smith     } else if (type == NORM_INFINITY) { /* max row norm */
1066e91c04dfSPierre Jolivet       PetscCall(MatNorm(mdn->A, type, nrm));
1067e91c04dfSPierre Jolivet       PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, nrm, 1, MPIU_REAL, MPIU_MAX, PetscObjectComm((PetscObject)A)));
1068ce94432eSBarry Smith     } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "No support for two norm");
10693501a2bdSLois Curfman McInnes   }
10709566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(mdn->A, &av));
10713ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
10723501a2bdSLois Curfman McInnes }
10733501a2bdSLois Curfman McInnes 
107466976f2fSJacob Faibussowitsch static PetscErrorCode MatTranspose_MPIDense(Mat A, MatReuse reuse, Mat *matout)
1075d71ae5a4SJacob Faibussowitsch {
10763501a2bdSLois Curfman McInnes   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
10773501a2bdSLois Curfman McInnes   Mat           B;
1078d0f46423SBarry Smith   PetscInt      M = A->rmap->N, N = A->cmap->N, m, n, *rwork, rstart = A->rmap->rstart;
1079637a0070SStefano Zampini   PetscInt      j, i, lda;
108087828ca2SBarry Smith   PetscScalar  *v;
10813501a2bdSLois Curfman McInnes 
10823a40ed3dSBarry Smith   PetscFunctionBegin;
10837fb60732SBarry Smith   if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(A, *matout));
1084cf37664fSBarry Smith   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX) {
10859566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
10869566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(B, A->cmap->n, A->rmap->n, N, M));
10879566063dSJacob Faibussowitsch     PetscCall(MatSetType(B, ((PetscObject)A)->type_name));
10889566063dSJacob Faibussowitsch     PetscCall(MatMPIDenseSetPreallocation(B, NULL));
1089637a0070SStefano Zampini   } else B = *matout;
10903501a2bdSLois Curfman McInnes 
10919371c9d4SSatish Balay   m = a->A->rmap->n;
10929371c9d4SSatish Balay   n = a->A->cmap->n;
10939566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)&v));
10949566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
10959566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(m, &rwork));
10963501a2bdSLois Curfman McInnes   for (i = 0; i < m; i++) rwork[i] = rstart + i;
10971acff37aSSatish Balay   for (j = 0; j < n; j++) {
10989566063dSJacob Faibussowitsch     PetscCall(MatSetValues(B, 1, &j, m, rwork, v, INSERT_VALUES));
10998e3a54c0SPierre Jolivet     v = PetscSafePointerPlusOffset(v, lda);
11003501a2bdSLois Curfman McInnes   }
11019566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)&v));
11029566063dSJacob Faibussowitsch   PetscCall(PetscFree(rwork));
11039566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
11049566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
1105cf37664fSBarry Smith   if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_REUSE_MATRIX) {
11063501a2bdSLois Curfman McInnes     *matout = B;
11073501a2bdSLois Curfman McInnes   } else {
11089566063dSJacob Faibussowitsch     PetscCall(MatHeaderMerge(A, &B));
11093501a2bdSLois Curfman McInnes   }
11103ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1111096963f5SLois Curfman McInnes }
1112096963f5SLois Curfman McInnes 
11136849ba73SBarry Smith static PetscErrorCode       MatDuplicate_MPIDense(Mat, MatDuplicateOption, Mat *);
111452c5f739Sprj- PETSC_INTERN PetscErrorCode MatScale_MPIDense(Mat, PetscScalar);
11158965ea79SLois Curfman McInnes 
111666976f2fSJacob Faibussowitsch static PetscErrorCode MatSetUp_MPIDense(Mat A)
1117d71ae5a4SJacob Faibussowitsch {
1118273d9f13SBarry Smith   PetscFunctionBegin;
11199566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(A->rmap));
11209566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(A->cmap));
112148a46eb9SPierre Jolivet   if (!A->preallocated) PetscCall(MatMPIDenseSetPreallocation(A, NULL));
11223ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1123273d9f13SBarry Smith }
1124273d9f13SBarry Smith 
112566976f2fSJacob Faibussowitsch static PetscErrorCode MatAXPY_MPIDense(Mat Y, PetscScalar alpha, Mat X, MatStructure str)
1126d71ae5a4SJacob Faibussowitsch {
1127488007eeSBarry Smith   Mat_MPIDense *A = (Mat_MPIDense *)Y->data, *B = (Mat_MPIDense *)X->data;
1128488007eeSBarry Smith 
1129488007eeSBarry Smith   PetscFunctionBegin;
11309566063dSJacob Faibussowitsch   PetscCall(MatAXPY(A->A, alpha, B->A, str));
11313ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1132488007eeSBarry Smith }
1133488007eeSBarry Smith 
113466976f2fSJacob Faibussowitsch static PetscErrorCode MatConjugate_MPIDense(Mat mat)
1135d71ae5a4SJacob Faibussowitsch {
1136ba337c44SJed Brown   Mat_MPIDense *a = (Mat_MPIDense *)mat->data;
1137ba337c44SJed Brown 
1138ba337c44SJed Brown   PetscFunctionBegin;
11399566063dSJacob Faibussowitsch   PetscCall(MatConjugate(a->A));
11403ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1141ba337c44SJed Brown }
1142ba337c44SJed Brown 
114366976f2fSJacob Faibussowitsch static PetscErrorCode MatRealPart_MPIDense(Mat A)
1144d71ae5a4SJacob Faibussowitsch {
1145ba337c44SJed Brown   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1146ba337c44SJed Brown 
1147ba337c44SJed Brown   PetscFunctionBegin;
11489566063dSJacob Faibussowitsch   PetscCall(MatRealPart(a->A));
11493ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1150ba337c44SJed Brown }
1151ba337c44SJed Brown 
115266976f2fSJacob Faibussowitsch static PetscErrorCode MatImaginaryPart_MPIDense(Mat A)
1153d71ae5a4SJacob Faibussowitsch {
1154ba337c44SJed Brown   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1155ba337c44SJed Brown 
1156ba337c44SJed Brown   PetscFunctionBegin;
11579566063dSJacob Faibussowitsch   PetscCall(MatImaginaryPart(a->A));
11583ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1159ba337c44SJed Brown }
1160ba337c44SJed Brown 
1161d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatGetColumnVector_MPIDense(Mat A, Vec v, PetscInt col)
1162d71ae5a4SJacob Faibussowitsch {
1163637a0070SStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
116449a6ff4bSBarry Smith 
116549a6ff4bSBarry Smith   PetscFunctionBegin;
116628b400f6SJacob Faibussowitsch   PetscCheck(a->A, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing local matrix");
116728b400f6SJacob Faibussowitsch   PetscCheck(a->A->ops->getcolumnvector, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing get column operation");
11689566063dSJacob Faibussowitsch   PetscCall((*a->A->ops->getcolumnvector)(a->A, v, col));
11693ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
117049a6ff4bSBarry Smith }
117149a6ff4bSBarry Smith 
1172857cbf51SRichard Tran Mills PETSC_INTERN PetscErrorCode MatGetColumnReductions_SeqDense(Mat, PetscInt, PetscReal *);
117352c5f739Sprj- 
117466976f2fSJacob Faibussowitsch static PetscErrorCode MatGetColumnReductions_MPIDense(Mat A, PetscInt type, PetscReal *reductions)
1175d71ae5a4SJacob Faibussowitsch {
1176a873a8cdSSam Reynolds   PetscInt      i, m, n;
11770716a85fSBarry Smith   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
11780716a85fSBarry Smith 
11790716a85fSBarry Smith   PetscFunctionBegin;
11809566063dSJacob Faibussowitsch   PetscCall(MatGetSize(A, &m, &n));
1181857cbf51SRichard Tran Mills   if (type == REDUCTION_MEAN_REALPART) {
1182e91c04dfSPierre Jolivet     PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_REALPART, reductions));
1183857cbf51SRichard Tran Mills   } else if (type == REDUCTION_MEAN_IMAGINARYPART) {
1184e91c04dfSPierre Jolivet     PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_IMAGINARYPART, reductions));
1185a873a8cdSSam Reynolds   } else {
1186e91c04dfSPierre Jolivet     PetscCall(MatGetColumnReductions_SeqDense(a->A, type, reductions));
1187a873a8cdSSam Reynolds   }
1188857cbf51SRichard Tran Mills   if (type == NORM_2) {
1189e91c04dfSPierre Jolivet     for (i = 0; i < n; i++) reductions[i] *= reductions[i];
11900716a85fSBarry Smith   }
1191857cbf51SRichard Tran Mills   if (type == NORM_INFINITY) {
1192e91c04dfSPierre Jolivet     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, reductions, n, MPIU_REAL, MPIU_MAX, A->hdr.comm));
11930716a85fSBarry Smith   } else {
1194e91c04dfSPierre Jolivet     PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, reductions, n, MPIU_REAL, MPIU_SUM, A->hdr.comm));
11950716a85fSBarry Smith   }
1196857cbf51SRichard Tran Mills   if (type == NORM_2) {
1197a873a8cdSSam Reynolds     for (i = 0; i < n; i++) reductions[i] = PetscSqrtReal(reductions[i]);
1198857cbf51SRichard Tran Mills   } else if (type == REDUCTION_MEAN_REALPART || type == REDUCTION_MEAN_IMAGINARYPART) {
1199a873a8cdSSam Reynolds     for (i = 0; i < n; i++) reductions[i] /= m;
12000716a85fSBarry Smith   }
12013ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
12020716a85fSBarry Smith }
12030716a85fSBarry Smith 
1204d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatSetRandom_MPIDense(Mat x, PetscRandom rctx)
1205d71ae5a4SJacob Faibussowitsch {
120673a71a0fSBarry Smith   Mat_MPIDense *d = (Mat_MPIDense *)x->data;
120773a71a0fSBarry Smith 
120873a71a0fSBarry Smith   PetscFunctionBegin;
12099566063dSJacob Faibussowitsch   PetscCall(MatSetRandom(d->A, rctx));
12103faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE)
12113faff063SStefano Zampini   x->offloadmask = d->A->offloadmask;
12123faff063SStefano Zampini #endif
12133ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
121473a71a0fSBarry Smith }
121573a71a0fSBarry Smith 
1216d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMissingDiagonal_MPIDense(Mat A, PetscBool *missing, PetscInt *d)
1217d71ae5a4SJacob Faibussowitsch {
12183b49f96aSBarry Smith   PetscFunctionBegin;
12193b49f96aSBarry Smith   *missing = PETSC_FALSE;
12203ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
12213b49f96aSBarry Smith }
12223b49f96aSBarry Smith 
12234222ddf1SHong Zhang static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat);
1224cc48ffa7SToby Isaac static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat);
12256718818eSStefano Zampini static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat);
1226a1f9a5e6SStefano Zampini static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat);
12276718818eSStefano Zampini static PetscErrorCode MatEqual_MPIDense(Mat, Mat, PetscBool *);
12286718818eSStefano Zampini static PetscErrorCode MatLoad_MPIDense(Mat, PetscViewer);
12295d8c7819SPierre Jolivet static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat);
1230cc48ffa7SToby Isaac 
123109dc0095SBarry Smith static struct _MatOps MatOps_Values = {MatSetValues_MPIDense,
123209dc0095SBarry Smith                                        MatGetRow_MPIDense,
123309dc0095SBarry Smith                                        MatRestoreRow_MPIDense,
123409dc0095SBarry Smith                                        MatMult_MPIDense,
123597304618SKris Buschelman                                        /*  4*/ MatMultAdd_MPIDense,
12367c922b88SBarry Smith                                        MatMultTranspose_MPIDense,
12377c922b88SBarry Smith                                        MatMultTransposeAdd_MPIDense,
1238f4259b30SLisandro Dalcin                                        NULL,
1239f4259b30SLisandro Dalcin                                        NULL,
1240f4259b30SLisandro Dalcin                                        NULL,
1241f4259b30SLisandro Dalcin                                        /* 10*/ NULL,
1242f4259b30SLisandro Dalcin                                        NULL,
1243f4259b30SLisandro Dalcin                                        NULL,
1244f4259b30SLisandro Dalcin                                        NULL,
124509dc0095SBarry Smith                                        MatTranspose_MPIDense,
124697304618SKris Buschelman                                        /* 15*/ MatGetInfo_MPIDense,
12476e4ee0c6SHong Zhang                                        MatEqual_MPIDense,
124809dc0095SBarry Smith                                        MatGetDiagonal_MPIDense,
12495b2fa520SLois Curfman McInnes                                        MatDiagonalScale_MPIDense,
125009dc0095SBarry Smith                                        MatNorm_MPIDense,
125197304618SKris Buschelman                                        /* 20*/ MatAssemblyBegin_MPIDense,
125209dc0095SBarry Smith                                        MatAssemblyEnd_MPIDense,
125309dc0095SBarry Smith                                        MatSetOption_MPIDense,
125409dc0095SBarry Smith                                        MatZeroEntries_MPIDense,
1255d519adbfSMatthew Knepley                                        /* 24*/ MatZeroRows_MPIDense,
1256f4259b30SLisandro Dalcin                                        NULL,
1257f4259b30SLisandro Dalcin                                        NULL,
1258f4259b30SLisandro Dalcin                                        NULL,
1259f4259b30SLisandro Dalcin                                        NULL,
12604994cf47SJed Brown                                        /* 29*/ MatSetUp_MPIDense,
1261f4259b30SLisandro Dalcin                                        NULL,
1262f4259b30SLisandro Dalcin                                        NULL,
1263c56a70eeSBarry Smith                                        MatGetDiagonalBlock_MPIDense,
1264f4259b30SLisandro Dalcin                                        NULL,
1265d519adbfSMatthew Knepley                                        /* 34*/ MatDuplicate_MPIDense,
1266f4259b30SLisandro Dalcin                                        NULL,
1267f4259b30SLisandro Dalcin                                        NULL,
1268f4259b30SLisandro Dalcin                                        NULL,
1269f4259b30SLisandro Dalcin                                        NULL,
1270d519adbfSMatthew Knepley                                        /* 39*/ MatAXPY_MPIDense,
12717dae84e0SHong Zhang                                        MatCreateSubMatrices_MPIDense,
1272f4259b30SLisandro Dalcin                                        NULL,
127309dc0095SBarry Smith                                        MatGetValues_MPIDense,
1274a76f77c3SStefano Zampini                                        MatCopy_MPIDense,
1275f4259b30SLisandro Dalcin                                        /* 44*/ NULL,
127609dc0095SBarry Smith                                        MatScale_MPIDense,
12772f605a99SJose E. Roman                                        MatShift_MPIDense,
1278f4259b30SLisandro Dalcin                                        NULL,
1279f4259b30SLisandro Dalcin                                        NULL,
128073a71a0fSBarry Smith                                        /* 49*/ MatSetRandom_MPIDense,
1281f4259b30SLisandro Dalcin                                        NULL,
1282f4259b30SLisandro Dalcin                                        NULL,
1283f4259b30SLisandro Dalcin                                        NULL,
1284f4259b30SLisandro Dalcin                                        NULL,
1285f4259b30SLisandro Dalcin                                        /* 54*/ NULL,
1286f4259b30SLisandro Dalcin                                        NULL,
1287f4259b30SLisandro Dalcin                                        NULL,
1288f4259b30SLisandro Dalcin                                        NULL,
1289f4259b30SLisandro Dalcin                                        NULL,
12907dae84e0SHong Zhang                                        /* 59*/ MatCreateSubMatrix_MPIDense,
1291b9b97703SBarry Smith                                        MatDestroy_MPIDense,
1292b9b97703SBarry Smith                                        MatView_MPIDense,
1293f4259b30SLisandro Dalcin                                        NULL,
1294f4259b30SLisandro Dalcin                                        NULL,
1295f4259b30SLisandro Dalcin                                        /* 64*/ NULL,
1296f4259b30SLisandro Dalcin                                        NULL,
1297f4259b30SLisandro Dalcin                                        NULL,
1298f4259b30SLisandro Dalcin                                        NULL,
1299f4259b30SLisandro Dalcin                                        NULL,
1300f4259b30SLisandro Dalcin                                        /* 69*/ NULL,
1301f4259b30SLisandro Dalcin                                        NULL,
1302f4259b30SLisandro Dalcin                                        NULL,
1303f4259b30SLisandro Dalcin                                        NULL,
1304f4259b30SLisandro Dalcin                                        NULL,
1305f4259b30SLisandro Dalcin                                        /* 74*/ NULL,
1306f4259b30SLisandro Dalcin                                        NULL,
1307f4259b30SLisandro Dalcin                                        NULL,
1308f4259b30SLisandro Dalcin                                        NULL,
1309f4259b30SLisandro Dalcin                                        NULL,
1310f4259b30SLisandro Dalcin                                        /* 79*/ NULL,
1311f4259b30SLisandro Dalcin                                        NULL,
1312f4259b30SLisandro Dalcin                                        NULL,
1313f4259b30SLisandro Dalcin                                        NULL,
13145bba2384SShri Abhyankar                                        /* 83*/ MatLoad_MPIDense,
1315f4259b30SLisandro Dalcin                                        NULL,
1316f4259b30SLisandro Dalcin                                        NULL,
1317f4259b30SLisandro Dalcin                                        NULL,
1318f4259b30SLisandro Dalcin                                        NULL,
1319f4259b30SLisandro Dalcin                                        NULL,
1320f4259b30SLisandro Dalcin                                        /* 89*/ NULL,
1321f4259b30SLisandro Dalcin                                        NULL,
1322f4259b30SLisandro Dalcin                                        NULL,
1323f4259b30SLisandro Dalcin                                        NULL,
1324f4259b30SLisandro Dalcin                                        NULL,
1325f4259b30SLisandro Dalcin                                        /* 94*/ NULL,
1326f4259b30SLisandro Dalcin                                        NULL,
13276718818eSStefano Zampini                                        MatMatTransposeMultSymbolic_MPIDense_MPIDense,
1328cc48ffa7SToby Isaac                                        MatMatTransposeMultNumeric_MPIDense_MPIDense,
1329f4259b30SLisandro Dalcin                                        NULL,
13304222ddf1SHong Zhang                                        /* 99*/ MatProductSetFromOptions_MPIDense,
1331f4259b30SLisandro Dalcin                                        NULL,
1332f4259b30SLisandro Dalcin                                        NULL,
1333ba337c44SJed Brown                                        MatConjugate_MPIDense,
1334f4259b30SLisandro Dalcin                                        NULL,
1335f4259b30SLisandro Dalcin                                        /*104*/ NULL,
1336ba337c44SJed Brown                                        MatRealPart_MPIDense,
1337ba337c44SJed Brown                                        MatImaginaryPart_MPIDense,
1338f4259b30SLisandro Dalcin                                        NULL,
1339f4259b30SLisandro Dalcin                                        NULL,
1340f4259b30SLisandro Dalcin                                        /*109*/ NULL,
1341f4259b30SLisandro Dalcin                                        NULL,
1342f4259b30SLisandro Dalcin                                        NULL,
134349a6ff4bSBarry Smith                                        MatGetColumnVector_MPIDense,
13443b49f96aSBarry Smith                                        MatMissingDiagonal_MPIDense,
1345f4259b30SLisandro Dalcin                                        /*114*/ NULL,
1346f4259b30SLisandro Dalcin                                        NULL,
1347f4259b30SLisandro Dalcin                                        NULL,
1348f4259b30SLisandro Dalcin                                        NULL,
1349f4259b30SLisandro Dalcin                                        NULL,
1350f4259b30SLisandro Dalcin                                        /*119*/ NULL,
1351f4259b30SLisandro Dalcin                                        NULL,
1352459e8d23SBlanca Mellado Pinto                                        MatMultHermitianTranspose_MPIDense,
1353459e8d23SBlanca Mellado Pinto                                        MatMultHermitianTransposeAdd_MPIDense,
1354f4259b30SLisandro Dalcin                                        NULL,
1355f4259b30SLisandro Dalcin                                        /*124*/ NULL,
1356a873a8cdSSam Reynolds                                        MatGetColumnReductions_MPIDense,
1357f4259b30SLisandro Dalcin                                        NULL,
1358f4259b30SLisandro Dalcin                                        NULL,
1359f4259b30SLisandro Dalcin                                        NULL,
1360f4259b30SLisandro Dalcin                                        /*129*/ NULL,
1361f4259b30SLisandro Dalcin                                        NULL,
13626718818eSStefano Zampini                                        MatTransposeMatMultSymbolic_MPIDense_MPIDense,
1363cb20be35SHong Zhang                                        MatTransposeMatMultNumeric_MPIDense_MPIDense,
1364f4259b30SLisandro Dalcin                                        NULL,
1365f4259b30SLisandro Dalcin                                        /*134*/ NULL,
1366f4259b30SLisandro Dalcin                                        NULL,
1367f4259b30SLisandro Dalcin                                        NULL,
1368f4259b30SLisandro Dalcin                                        NULL,
1369f4259b30SLisandro Dalcin                                        NULL,
1370f4259b30SLisandro Dalcin                                        /*139*/ NULL,
1371f4259b30SLisandro Dalcin                                        NULL,
1372f4259b30SLisandro Dalcin                                        NULL,
1373f4259b30SLisandro Dalcin                                        NULL,
1374f4259b30SLisandro Dalcin                                        NULL,
13754222ddf1SHong Zhang                                        MatCreateMPIMatConcatenateSeqMat_MPIDense,
1376f4259b30SLisandro Dalcin                                        /*145*/ NULL,
1377f4259b30SLisandro Dalcin                                        NULL,
137899a7f59eSMark Adams                                        NULL,
137999a7f59eSMark Adams                                        NULL,
13807fb60732SBarry Smith                                        NULL,
1381dec0b466SHong Zhang                                        /*150*/ NULL,
1382eede4a3fSMark Adams                                        NULL,
13834cc2b5b5SPierre Jolivet                                        NULL,
138442ce410bSJunchao Zhang                                        NULL,
138542ce410bSJunchao Zhang                                        NULL,
1386fe1fc275SAlexander                                        /*155*/ NULL,
1387dec0b466SHong Zhang                                        NULL};
13888965ea79SLois Curfman McInnes 
138966976f2fSJacob Faibussowitsch static PetscErrorCode MatMPIDenseSetPreallocation_MPIDense(Mat mat, PetscScalar *data)
1390d71ae5a4SJacob Faibussowitsch {
1391637a0070SStefano Zampini   Mat_MPIDense *a     = (Mat_MPIDense *)mat->data;
139247d993e7Ssuyashtn   MatType       mtype = MATSEQDENSE;
1393a23d5eceSKris Buschelman 
1394a23d5eceSKris Buschelman   PetscFunctionBegin;
139528b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
13969566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(mat->rmap));
13979566063dSJacob Faibussowitsch   PetscCall(PetscLayoutSetUp(mat->cmap));
1398637a0070SStefano Zampini   if (!a->A) {
13999566063dSJacob Faibussowitsch     PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
14009566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(a->A, mat->rmap->n, mat->cmap->N, mat->rmap->n, mat->cmap->N));
1401637a0070SStefano Zampini   }
140247d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA)
140347d993e7Ssuyashtn   PetscBool iscuda;
14049566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSECUDA, &iscuda));
140547d993e7Ssuyashtn   if (iscuda) mtype = MATSEQDENSECUDA;
140647d993e7Ssuyashtn #endif
140747d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
140847d993e7Ssuyashtn   PetscBool iship;
140947d993e7Ssuyashtn   PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSEHIP, &iship));
141047d993e7Ssuyashtn   if (iship) mtype = MATSEQDENSEHIP;
141147d993e7Ssuyashtn #endif
141247d993e7Ssuyashtn   PetscCall(MatSetType(a->A, mtype));
14139566063dSJacob Faibussowitsch   PetscCall(MatSeqDenseSetPreallocation(a->A, data));
141447d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
14153faff063SStefano Zampini   mat->offloadmask = a->A->offloadmask;
14163faff063SStefano Zampini #endif
1417637a0070SStefano Zampini   mat->preallocated = PETSC_TRUE;
14186f08ad05SPierre Jolivet   mat->assembled    = PETSC_TRUE;
14193ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1420a23d5eceSKris Buschelman }
1421a23d5eceSKris Buschelman 
1422d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode MatConvert_MPIAIJ_MPIDense(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1423d71ae5a4SJacob Faibussowitsch {
14244e29119aSPierre Jolivet   Mat B, C;
14254e29119aSPierre Jolivet 
14264e29119aSPierre Jolivet   PetscFunctionBegin;
14279566063dSJacob Faibussowitsch   PetscCall(MatMPIAIJGetLocalMat(A, MAT_INITIAL_MATRIX, &C));
14289566063dSJacob Faibussowitsch   PetscCall(MatConvert_SeqAIJ_SeqDense(C, MATSEQDENSE, MAT_INITIAL_MATRIX, &B));
14299566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&C));
14304e29119aSPierre Jolivet   if (reuse == MAT_REUSE_MATRIX) {
14314e29119aSPierre Jolivet     C = *newmat;
14324e29119aSPierre Jolivet   } else C = NULL;
14339566063dSJacob Faibussowitsch   PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C));
14349566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&B));
14354e29119aSPierre Jolivet   if (reuse == MAT_INPLACE_MATRIX) {
14369566063dSJacob Faibussowitsch     PetscCall(MatHeaderReplace(A, &C));
14374e29119aSPierre Jolivet   } else if (reuse == MAT_INITIAL_MATRIX) *newmat = C;
14383ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14394e29119aSPierre Jolivet }
14404e29119aSPierre Jolivet 
144166976f2fSJacob Faibussowitsch static PetscErrorCode MatConvert_MPIDense_MPIAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1442d71ae5a4SJacob Faibussowitsch {
14434e29119aSPierre Jolivet   Mat B, C;
14444e29119aSPierre Jolivet 
14454e29119aSPierre Jolivet   PetscFunctionBegin;
14469566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLocalMatrix(A, &C));
14479566063dSJacob Faibussowitsch   PetscCall(MatConvert_SeqDense_SeqAIJ(C, MATSEQAIJ, MAT_INITIAL_MATRIX, &B));
14484e29119aSPierre Jolivet   if (reuse == MAT_REUSE_MATRIX) {
14494e29119aSPierre Jolivet     C = *newmat;
14504e29119aSPierre Jolivet   } else C = NULL;
14519566063dSJacob Faibussowitsch   PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C));
14529566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&B));
14534e29119aSPierre Jolivet   if (reuse == MAT_INPLACE_MATRIX) {
14549566063dSJacob Faibussowitsch     PetscCall(MatHeaderReplace(A, &C));
14554e29119aSPierre Jolivet   } else if (reuse == MAT_INITIAL_MATRIX) *newmat = C;
14563ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
14574e29119aSPierre Jolivet }
14584e29119aSPierre Jolivet 
145965b80a83SHong Zhang #if defined(PETSC_HAVE_ELEMENTAL)
1460d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode MatConvert_MPIDense_Elemental(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1461d71ae5a4SJacob Faibussowitsch {
1462f709e8d7SPierre Jolivet   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
14638ea901baSHong Zhang   Mat           mat_elemental;
146432d7a744SHong Zhang   PetscScalar  *v;
1465f709e8d7SPierre Jolivet   PetscInt      m = A->rmap->n, N = A->cmap->N, rstart = A->rmap->rstart, i, *rows, *cols, lda;
14668ea901baSHong Zhang 
14678baccfbdSHong Zhang   PetscFunctionBegin;
1468378336b6SHong Zhang   if (reuse == MAT_REUSE_MATRIX) {
1469378336b6SHong Zhang     mat_elemental = *newmat;
14709566063dSJacob Faibussowitsch     PetscCall(MatZeroEntries(*newmat));
1471378336b6SHong Zhang   } else {
14729566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat_elemental));
14739566063dSJacob Faibussowitsch     PetscCall(MatSetSizes(mat_elemental, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N));
14749566063dSJacob Faibussowitsch     PetscCall(MatSetType(mat_elemental, MATELEMENTAL));
14759566063dSJacob Faibussowitsch     PetscCall(MatSetUp(mat_elemental));
14769566063dSJacob Faibussowitsch     PetscCall(MatSetOption(mat_elemental, MAT_ROW_ORIENTED, PETSC_FALSE));
1477378336b6SHong Zhang   }
1478378336b6SHong Zhang 
14799566063dSJacob Faibussowitsch   PetscCall(PetscMalloc2(m, &rows, N, &cols));
148032d7a744SHong Zhang   for (i = 0; i < N; i++) cols[i] = i;
148132d7a744SHong Zhang   for (i = 0; i < m; i++) rows[i] = rstart + i;
14828ea901baSHong Zhang 
1483637a0070SStefano Zampini   /* PETSc-Elemental interface uses axpy for setting off-processor entries, only ADD_VALUES is allowed */
14849566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(A, &v));
1485f709e8d7SPierre Jolivet   PetscCall(MatDenseGetLDA(a->A, &lda));
1486f709e8d7SPierre Jolivet   if (lda == m) PetscCall(MatSetValues(mat_elemental, m, rows, N, cols, v, ADD_VALUES));
1487f709e8d7SPierre Jolivet   else {
1488f709e8d7SPierre Jolivet     for (i = 0; i < N; i++) PetscCall(MatSetValues(mat_elemental, m, rows, 1, &i, v + lda * i, ADD_VALUES));
1489f709e8d7SPierre Jolivet   }
14909566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(mat_elemental, MAT_FINAL_ASSEMBLY));
14919566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(mat_elemental, MAT_FINAL_ASSEMBLY));
14929566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(A, &v));
14939566063dSJacob Faibussowitsch   PetscCall(PetscFree2(rows, cols));
14948ea901baSHong Zhang 
1495511c6705SHong Zhang   if (reuse == MAT_INPLACE_MATRIX) {
14969566063dSJacob Faibussowitsch     PetscCall(MatHeaderReplace(A, &mat_elemental));
14978ea901baSHong Zhang   } else {
14988ea901baSHong Zhang     *newmat = mat_elemental;
14998ea901baSHong Zhang   }
15003ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15018baccfbdSHong Zhang }
150265b80a83SHong Zhang #endif
15038baccfbdSHong Zhang 
1504d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetColumn_MPIDense(Mat A, PetscInt col, PetscScalar **vals)
1505d71ae5a4SJacob Faibussowitsch {
150686aefd0dSHong Zhang   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
150786aefd0dSHong Zhang 
150886aefd0dSHong Zhang   PetscFunctionBegin;
15099566063dSJacob Faibussowitsch   PetscCall(MatDenseGetColumn(mat->A, col, vals));
15103ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
151186aefd0dSHong Zhang }
151286aefd0dSHong Zhang 
1513d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseRestoreColumn_MPIDense(Mat A, PetscScalar **vals)
1514d71ae5a4SJacob Faibussowitsch {
151586aefd0dSHong Zhang   Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
151686aefd0dSHong Zhang 
151786aefd0dSHong Zhang   PetscFunctionBegin;
15189566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreColumn(mat->A, vals));
15193ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
152086aefd0dSHong Zhang }
152186aefd0dSHong Zhang 
1522d71ae5a4SJacob Faibussowitsch PetscErrorCode MatCreateMPIMatConcatenateSeqMat_MPIDense(MPI_Comm comm, Mat inmat, PetscInt n, MatReuse scall, Mat *outmat)
1523d71ae5a4SJacob Faibussowitsch {
152494e2cb23SJakub Kruzik   Mat_MPIDense *mat;
152594e2cb23SJakub Kruzik   PetscInt      m, nloc, N;
152694e2cb23SJakub Kruzik 
152794e2cb23SJakub Kruzik   PetscFunctionBegin;
15289566063dSJacob Faibussowitsch   PetscCall(MatGetSize(inmat, &m, &N));
15299566063dSJacob Faibussowitsch   PetscCall(MatGetLocalSize(inmat, NULL, &nloc));
153094e2cb23SJakub Kruzik   if (scall == MAT_INITIAL_MATRIX) { /* symbolic phase */
153194e2cb23SJakub Kruzik     PetscInt sum;
153294e2cb23SJakub Kruzik 
153348a46eb9SPierre Jolivet     if (n == PETSC_DECIDE) PetscCall(PetscSplitOwnership(comm, &n, &N));
153494e2cb23SJakub Kruzik     /* Check sum(n) = N */
1535462c564dSBarry Smith     PetscCallMPI(MPIU_Allreduce(&n, &sum, 1, MPIU_INT, MPI_SUM, comm));
153608401ef6SPierre Jolivet     PetscCheck(sum == N, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Sum of local columns %" PetscInt_FMT " != global columns %" PetscInt_FMT, sum, N);
153794e2cb23SJakub Kruzik 
15389566063dSJacob Faibussowitsch     PetscCall(MatCreateDense(comm, m, n, PETSC_DETERMINE, N, NULL, outmat));
15399566063dSJacob Faibussowitsch     PetscCall(MatSetOption(*outmat, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
154094e2cb23SJakub Kruzik   }
154194e2cb23SJakub Kruzik 
154294e2cb23SJakub Kruzik   /* numeric phase */
154394e2cb23SJakub Kruzik   mat = (Mat_MPIDense *)(*outmat)->data;
15449566063dSJacob Faibussowitsch   PetscCall(MatCopy(inmat, mat->A, SAME_NONZERO_PATTERN));
15453ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
154694e2cb23SJakub Kruzik }
154794e2cb23SJakub Kruzik 
1548d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVec_MPIDense(Mat A, PetscInt col, Vec *v)
1549d71ae5a4SJacob Faibussowitsch {
15506947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15516947451fSStefano Zampini   PetscInt      lda;
15526947451fSStefano Zampini 
15536947451fSStefano Zampini   PetscFunctionBegin;
155428b400f6SJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
155528b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1556d16ceb75SStefano Zampini   if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
15576947451fSStefano Zampini   a->vecinuse = col + 1;
15589566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
15599566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArray(a->A, (PetscScalar **)&a->ptrinuse));
15608e3a54c0SPierre Jolivet   PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
15616947451fSStefano Zampini   *v = a->cvec;
15623ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15636947451fSStefano Zampini }
15646947451fSStefano Zampini 
1565d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVec_MPIDense(Mat A, PetscInt col, Vec *v)
1566d71ae5a4SJacob Faibussowitsch {
15676947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15686947451fSStefano Zampini 
15696947451fSStefano Zampini   PetscFunctionBegin;
157028b400f6SJacob Faibussowitsch   PetscCheck(a->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
157128b400f6SJacob Faibussowitsch   PetscCheck(a->cvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column vector");
15724186a4bbSPierre Jolivet   VecCheckAssembled(a->cvec);
15736947451fSStefano Zampini   a->vecinuse = 0;
15749566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArray(a->A, (PetscScalar **)&a->ptrinuse));
15759566063dSJacob Faibussowitsch   PetscCall(VecResetArray(a->cvec));
1576742765d3SMatthew Knepley   if (v) *v = NULL;
15773ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15786947451fSStefano Zampini }
15796947451fSStefano Zampini 
1580d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v)
1581d71ae5a4SJacob Faibussowitsch {
15826947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15836947451fSStefano Zampini   PetscInt      lda;
15846947451fSStefano Zampini 
15856947451fSStefano Zampini   PetscFunctionBegin;
158628b400f6SJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
158728b400f6SJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1588d16ceb75SStefano Zampini   if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
15896947451fSStefano Zampini   a->vecinuse = col + 1;
15909566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
15919566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, &a->ptrinuse));
15928e3a54c0SPierre Jolivet   PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
15939566063dSJacob Faibussowitsch   PetscCall(VecLockReadPush(a->cvec));
15946947451fSStefano Zampini   *v = a->cvec;
15953ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
15966947451fSStefano Zampini }
15976947451fSStefano Zampini 
1598d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v)
1599d71ae5a4SJacob Faibussowitsch {
16006947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
16016947451fSStefano Zampini 
16026947451fSStefano Zampini   PetscFunctionBegin;
16035f80ce2aSJacob Faibussowitsch   PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
16045f80ce2aSJacob Faibussowitsch   PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector");
16054186a4bbSPierre Jolivet   VecCheckAssembled(a->cvec);
16066947451fSStefano Zampini   a->vecinuse = 0;
16079566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, &a->ptrinuse));
16089566063dSJacob Faibussowitsch   PetscCall(VecLockReadPop(a->cvec));
16099566063dSJacob Faibussowitsch   PetscCall(VecResetArray(a->cvec));
1610742765d3SMatthew Knepley   if (v) *v = NULL;
16113ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
16126947451fSStefano Zampini }
16136947451fSStefano Zampini 
1614d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v)
1615d71ae5a4SJacob Faibussowitsch {
16166947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
16176947451fSStefano Zampini   PetscInt      lda;
16186947451fSStefano Zampini 
16196947451fSStefano Zampini   PetscFunctionBegin;
16205f80ce2aSJacob Faibussowitsch   PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
16215f80ce2aSJacob Faibussowitsch   PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1622d16ceb75SStefano Zampini   if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
16236947451fSStefano Zampini   a->vecinuse = col + 1;
16249566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &lda));
16259566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(a->A, (PetscScalar **)&a->ptrinuse));
16268e3a54c0SPierre Jolivet   PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
16276947451fSStefano Zampini   *v = a->cvec;
16283ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
16296947451fSStefano Zampini }
16306947451fSStefano Zampini 
1631d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v)
1632d71ae5a4SJacob Faibussowitsch {
16336947451fSStefano Zampini   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
16346947451fSStefano Zampini 
16356947451fSStefano Zampini   PetscFunctionBegin;
16365f80ce2aSJacob Faibussowitsch   PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
16375f80ce2aSJacob Faibussowitsch   PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector");
16384186a4bbSPierre Jolivet   VecCheckAssembled(a->cvec);
16396947451fSStefano Zampini   a->vecinuse = 0;
16409566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(a->A, (PetscScalar **)&a->ptrinuse));
16419566063dSJacob Faibussowitsch   PetscCall(VecResetArray(a->cvec));
1642742765d3SMatthew Knepley   if (v) *v = NULL;
16433ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
16446947451fSStefano Zampini }
16456947451fSStefano Zampini 
164666976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseGetSubMatrix_MPIDense(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *v)
1647d71ae5a4SJacob Faibussowitsch {
16485ea7661aSPierre Jolivet   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1649616b8fbbSStefano Zampini   Mat_MPIDense *c;
1650616b8fbbSStefano Zampini   MPI_Comm      comm;
1651d016bddeSToby Isaac   PetscInt      prbegin, prend, pcbegin, pcend;
16525ea7661aSPierre Jolivet 
16535ea7661aSPierre Jolivet   PetscFunctionBegin;
16549566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
16555f80ce2aSJacob Faibussowitsch   PetscCheck(!a->vecinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
16565f80ce2aSJacob Faibussowitsch   PetscCheck(!a->matinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1657d016bddeSToby Isaac   prbegin = PetscMax(0, PetscMin(A->rmap->rend, rbegin) - A->rmap->rstart);
1658d016bddeSToby Isaac   prend   = PetscMin(A->rmap->n, PetscMax(0, rend - A->rmap->rstart));
1659d016bddeSToby Isaac   pcbegin = PetscMax(0, PetscMin(A->cmap->rend, cbegin) - A->cmap->rstart);
1660d016bddeSToby Isaac   pcend   = PetscMin(A->cmap->n, PetscMax(0, cend - A->cmap->rstart));
16615ea7661aSPierre Jolivet   if (!a->cmat) {
16629566063dSJacob Faibussowitsch     PetscCall(MatCreate(comm, &a->cmat));
16639566063dSJacob Faibussowitsch     PetscCall(MatSetType(a->cmat, ((PetscObject)A)->type_name));
1664a2748737SPierre Jolivet     if (rend - rbegin == A->rmap->N) PetscCall(PetscLayoutReference(A->rmap, &a->cmat->rmap));
1665a2748737SPierre Jolivet     else {
1666d016bddeSToby Isaac       PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, prend - prbegin));
1667a2748737SPierre Jolivet       PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin));
1668a2748737SPierre Jolivet       PetscCall(PetscLayoutSetUp(a->cmat->rmap));
1669a2748737SPierre Jolivet     }
1670d016bddeSToby Isaac     if (cend - cbegin == A->cmap->N) PetscCall(PetscLayoutReference(A->cmap, &a->cmat->cmap));
1671d016bddeSToby Isaac     else {
1672d016bddeSToby Isaac       PetscCall(PetscLayoutSetLocalSize(a->cmat->cmap, pcend - pcbegin));
16739566063dSJacob Faibussowitsch       PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin));
16749566063dSJacob Faibussowitsch       PetscCall(PetscLayoutSetUp(a->cmat->cmap));
1675a2748737SPierre Jolivet     }
1676616b8fbbSStefano Zampini     c             = (Mat_MPIDense *)a->cmat->data;
1677d016bddeSToby Isaac     c->sub_rbegin = rbegin;
1678d016bddeSToby Isaac     c->sub_rend   = rend;
1679d016bddeSToby Isaac     c->sub_cbegin = cbegin;
1680d016bddeSToby Isaac     c->sub_cend   = cend;
1681d016bddeSToby Isaac   }
1682d016bddeSToby Isaac   c = (Mat_MPIDense *)a->cmat->data;
1683d016bddeSToby Isaac   if (c->sub_rbegin != rbegin || c->sub_rend != rend) {
1684d016bddeSToby Isaac     PetscCall(PetscLayoutDestroy(&a->cmat->rmap));
1685d016bddeSToby Isaac     PetscCall(PetscLayoutCreate(comm, &a->cmat->rmap));
1686d016bddeSToby Isaac     PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, prend - prbegin));
1687d016bddeSToby Isaac     PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin));
1688d016bddeSToby Isaac     PetscCall(PetscLayoutSetUp(a->cmat->rmap));
1689d016bddeSToby Isaac     c->sub_rbegin = rbegin;
1690d016bddeSToby Isaac     c->sub_rend   = rend;
1691d016bddeSToby Isaac   }
1692d016bddeSToby Isaac   if (c->sub_cbegin != cbegin || c->sub_cend != cend) {
1693d016bddeSToby Isaac     // special optimization: check if all columns are owned by rank 0, in which case no communication is necessary
1694d016bddeSToby Isaac     if ((cend - cbegin != a->cmat->cmap->N) || (A->cmap->range[1] != A->cmap->N)) {
1695d016bddeSToby Isaac       PetscCall(PetscLayoutDestroy(&a->cmat->cmap));
1696d016bddeSToby Isaac       PetscCall(PetscLayoutCreate(comm, &a->cmat->cmap));
1697d016bddeSToby Isaac       PetscCall(PetscLayoutSetLocalSize(a->cmat->cmap, pcend - pcbegin));
1698d016bddeSToby Isaac       PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin));
1699d016bddeSToby Isaac       PetscCall(PetscLayoutSetUp(a->cmat->cmap));
1700d016bddeSToby Isaac       PetscCall(VecDestroy(&c->lvec));
1701d016bddeSToby Isaac       PetscCall(PetscSFDestroy(&c->Mvctx));
1702d016bddeSToby Isaac     }
1703d016bddeSToby Isaac     c->sub_cbegin = cbegin;
1704d016bddeSToby Isaac     c->sub_cend   = cend;
1705d016bddeSToby Isaac   }
17065f80ce2aSJacob Faibussowitsch   PetscCheck(!c->A, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1707d016bddeSToby Isaac   PetscCall(MatDenseGetSubMatrix(a->A, prbegin, prend, cbegin, cend, &c->A));
17083faff063SStefano Zampini 
1709616b8fbbSStefano Zampini   a->cmat->preallocated = PETSC_TRUE;
1710616b8fbbSStefano Zampini   a->cmat->assembled    = PETSC_TRUE;
17113faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE)
17123faff063SStefano Zampini   a->cmat->offloadmask = c->A->offloadmask;
17133faff063SStefano Zampini #endif
17145ea7661aSPierre Jolivet   a->matinuse = cbegin + 1;
17155ea7661aSPierre Jolivet   *v          = a->cmat;
17163ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
17175ea7661aSPierre Jolivet }
17185ea7661aSPierre Jolivet 
171966976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreSubMatrix_MPIDense(Mat A, Mat *v)
1720d71ae5a4SJacob Faibussowitsch {
17215ea7661aSPierre Jolivet   Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1722616b8fbbSStefano Zampini   Mat_MPIDense *c;
17235ea7661aSPierre Jolivet 
17245ea7661aSPierre Jolivet   PetscFunctionBegin;
17255f80ce2aSJacob Faibussowitsch   PetscCheck(a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetSubMatrix() first");
17265f80ce2aSJacob Faibussowitsch   PetscCheck(a->cmat, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal matrix");
17275f80ce2aSJacob Faibussowitsch   PetscCheck(*v == a->cmat, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Not the matrix obtained from MatDenseGetSubMatrix()");
17285ea7661aSPierre Jolivet   a->matinuse = 0;
1729616b8fbbSStefano Zampini   c           = (Mat_MPIDense *)a->cmat->data;
17309566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreSubMatrix(a->A, &c->A));
1731742765d3SMatthew Knepley   if (v) *v = NULL;
17323faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE)
17333faff063SStefano Zampini   A->offloadmask = a->A->offloadmask;
17343faff063SStefano Zampini #endif
17353ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
17365ea7661aSPierre Jolivet }
17375ea7661aSPierre Jolivet 
1738002bc6daSRichard Tran Mills /*MC
1739002bc6daSRichard Tran Mills    MATMPIDENSE - MATMPIDENSE = "mpidense" - A matrix type to be used for distributed dense matrices.
1740002bc6daSRichard Tran Mills 
17412ef1f0ffSBarry Smith    Options Database Key:
174211a5261eSBarry Smith . -mat_type mpidense - sets the matrix type to `MATMPIDENSE` during a call to `MatSetFromOptions()`
1743002bc6daSRichard Tran Mills 
1744002bc6daSRichard Tran Mills   Level: beginner
1745002bc6daSRichard Tran Mills 
17461cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MatCreateDense()`, `MATSEQDENSE`, `MATDENSE`
1747002bc6daSRichard Tran Mills M*/
17484742e46bSJacob Faibussowitsch PetscErrorCode MatCreate_MPIDense(Mat mat)
1749d71ae5a4SJacob Faibussowitsch {
1750273d9f13SBarry Smith   Mat_MPIDense *a;
1751273d9f13SBarry Smith 
1752273d9f13SBarry Smith   PetscFunctionBegin;
17534dfa11a4SJacob Faibussowitsch   PetscCall(PetscNew(&a));
1754b0a32e0cSBarry Smith   mat->data   = (void *)a;
1755aea10558SJacob Faibussowitsch   mat->ops[0] = MatOps_Values;
1756273d9f13SBarry Smith 
1757273d9f13SBarry Smith   mat->insertmode = NOT_SET_VALUES;
1758273d9f13SBarry Smith 
1759273d9f13SBarry Smith   /* build cache for off array entries formed */
1760273d9f13SBarry Smith   a->donotstash = PETSC_FALSE;
17612205254eSKarl Rupp 
17629566063dSJacob Faibussowitsch   PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)mat), 1, &mat->stash));
1763273d9f13SBarry Smith 
1764273d9f13SBarry Smith   /* stuff used for matrix vector multiply */
1765f4259b30SLisandro Dalcin   a->lvec        = NULL;
1766f4259b30SLisandro Dalcin   a->Mvctx       = NULL;
1767273d9f13SBarry Smith   a->roworiented = PETSC_TRUE;
1768273d9f13SBarry Smith 
17699566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", MatDenseGetLDA_MPIDense));
17709566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", MatDenseSetLDA_MPIDense));
17719566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", MatDenseGetArray_MPIDense));
17729566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", MatDenseRestoreArray_MPIDense));
17739566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", MatDenseGetArrayRead_MPIDense));
17749566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", MatDenseRestoreArrayRead_MPIDense));
17759566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", MatDenseGetArrayWrite_MPIDense));
17769566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", MatDenseRestoreArrayWrite_MPIDense));
17779566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", MatDensePlaceArray_MPIDense));
17789566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", MatDenseResetArray_MPIDense));
17799566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", MatDenseReplaceArray_MPIDense));
17809566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_MPIDense));
17819566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_MPIDense));
17829566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_MPIDense));
17839566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_MPIDense));
17849566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_MPIDense));
17859566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_MPIDense));
17869566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_MPIDense));
17879566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_MPIDense));
17889566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", MatConvert_MPIAIJ_MPIDense));
17899566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", MatConvert_MPIDense_MPIAIJ));
17908baccfbdSHong Zhang #if defined(PETSC_HAVE_ELEMENTAL)
17919566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", MatConvert_MPIDense_Elemental));
17928baccfbdSHong Zhang #endif
1793d24d4204SJose E. Roman #if defined(PETSC_HAVE_SCALAPACK)
17949566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", MatConvert_Dense_ScaLAPACK));
1795d24d4204SJose E. Roman #endif
1796637a0070SStefano Zampini #if defined(PETSC_HAVE_CUDA)
17979566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", MatConvert_MPIDense_MPIDenseCUDA));
1798637a0070SStefano Zampini #endif
17999566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", MatMPIDenseSetPreallocation_MPIDense));
18009566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
18019566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
18026718818eSStefano Zampini #if defined(PETSC_HAVE_CUDA)
18039566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
18049566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
18056718818eSStefano Zampini #endif
180647d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
180747d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", MatConvert_MPIDense_MPIDenseHIP));
180847d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
180947d993e7Ssuyashtn   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
181047d993e7Ssuyashtn #endif
18119566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", MatDenseGetColumn_MPIDense));
18129566063dSJacob Faibussowitsch   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", MatDenseRestoreColumn_MPIDense));
1813d016bddeSToby Isaac   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultColumnRange_C", MatMultColumnRange_MPIDense));
18140be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", MatMultAddColumnRange_MPIDense));
18150be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", MatMultHermitianTransposeColumnRange_MPIDense));
18160be0d8bdSHansol Suh   PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", MatMultHermitianTransposeAddColumnRange_MPIDense));
18179566063dSJacob Faibussowitsch   PetscCall(PetscObjectChangeTypeName((PetscObject)mat, MATMPIDENSE));
18183ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1819273d9f13SBarry Smith }
1820273d9f13SBarry Smith 
1821209238afSKris Buschelman /*MC
1822002d173eSKris Buschelman    MATDENSE - MATDENSE = "dense" - A matrix type to be used for dense matrices.
1823209238afSKris Buschelman 
182411a5261eSBarry Smith    This matrix type is identical to `MATSEQDENSE` when constructed with a single process communicator,
182511a5261eSBarry Smith    and `MATMPIDENSE` otherwise.
1826209238afSKris Buschelman 
18272ef1f0ffSBarry Smith    Options Database Key:
182811a5261eSBarry Smith . -mat_type dense - sets the matrix type to `MATDENSE` during a call to `MatSetFromOptions()`
1829209238afSKris Buschelman 
1830209238afSKris Buschelman   Level: beginner
1831209238afSKris Buschelman 
18321cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATSEQDENSE`, `MATMPIDENSE`, `MATDENSECUDA`, `MATDENSEHIP`
18336947451fSStefano Zampini M*/
18346947451fSStefano Zampini 
18355d83a8b1SBarry Smith /*@
1836273d9f13SBarry Smith   MatMPIDenseSetPreallocation - Sets the array used to store the matrix entries
1837273d9f13SBarry Smith 
18388f490ea3SStefano Zampini   Collective
1839273d9f13SBarry Smith 
1840273d9f13SBarry Smith   Input Parameters:
1841fe59aa6dSJacob Faibussowitsch + B    - the matrix
18422ef1f0ffSBarry Smith - data - optional location of matrix data.  Set to `NULL` for PETSc
1843273d9f13SBarry Smith          to control all matrix memory allocation.
1844273d9f13SBarry Smith 
18452ef1f0ffSBarry Smith   Level: intermediate
18462ef1f0ffSBarry Smith 
1847273d9f13SBarry Smith   Notes:
18482ef1f0ffSBarry Smith   The dense format is fully compatible with standard Fortran
1849273d9f13SBarry Smith   storage by columns.
1850273d9f13SBarry Smith 
1851273d9f13SBarry Smith   The data input variable is intended primarily for Fortran programmers
1852273d9f13SBarry Smith   who wish to allocate their own matrix memory space.  Most users should
18532ef1f0ffSBarry Smith   set `data` to `NULL`.
1854273d9f13SBarry Smith 
18551cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATMPIDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()`
1856273d9f13SBarry Smith @*/
1857d71ae5a4SJacob Faibussowitsch PetscErrorCode MatMPIDenseSetPreallocation(Mat B, PetscScalar *data)
1858d71ae5a4SJacob Faibussowitsch {
1859273d9f13SBarry Smith   PetscFunctionBegin;
1860d5ea218eSStefano Zampini   PetscValidHeaderSpecific(B, MAT_CLASSID, 1);
1861cac4c232SBarry Smith   PetscTryMethod(B, "MatMPIDenseSetPreallocation_C", (Mat, PetscScalar *), (B, data));
18623ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1863273d9f13SBarry Smith }
1864273d9f13SBarry Smith 
1865d3042a70SBarry Smith /*@
186611a5261eSBarry Smith   MatDensePlaceArray - Allows one to replace the array in a `MATDENSE` matrix with an
1867d3042a70SBarry Smith   array provided by the user. This is useful to avoid copying an array
1868d3042a70SBarry Smith   into a matrix
1869d3042a70SBarry Smith 
1870d3042a70SBarry Smith   Not Collective
1871d3042a70SBarry Smith 
1872d3042a70SBarry Smith   Input Parameters:
1873d3042a70SBarry Smith + mat   - the matrix
1874d3042a70SBarry Smith - array - the array in column major order
1875d3042a70SBarry Smith 
18762ef1f0ffSBarry Smith   Level: developer
18772ef1f0ffSBarry Smith 
187811a5261eSBarry Smith   Note:
18794d5500e8SJunchao Zhang   Adding `const` to `array` was an oversight, see notes in `VecPlaceArray()`.
18804d5500e8SJunchao Zhang 
188111a5261eSBarry 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
1882d3042a70SBarry Smith   freed when the matrix is destroyed.
1883d3042a70SBarry Smith 
18841cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDenseResetArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`,
188511a5261eSBarry Smith           `MatDenseReplaceArray()`
1886d3042a70SBarry Smith @*/
1887d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDensePlaceArray(Mat mat, const PetscScalar *array)
1888d71ae5a4SJacob Faibussowitsch {
1889d3042a70SBarry Smith   PetscFunctionBegin;
1890d5ea218eSStefano Zampini   PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
1891cac4c232SBarry Smith   PetscUseMethod(mat, "MatDensePlaceArray_C", (Mat, const PetscScalar *), (mat, array));
18929566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)mat));
189347d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
1894637a0070SStefano Zampini   mat->offloadmask = PETSC_OFFLOAD_CPU;
1895637a0070SStefano Zampini #endif
18963ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1897d3042a70SBarry Smith }
1898d3042a70SBarry Smith 
1899d3042a70SBarry Smith /*@
190011a5261eSBarry Smith   MatDenseResetArray - Resets the matrix array to that it previously had before the call to `MatDensePlaceArray()`
1901d3042a70SBarry Smith 
1902d3042a70SBarry Smith   Not Collective
1903d3042a70SBarry Smith 
19042fe279fdSBarry Smith   Input Parameter:
1905d3042a70SBarry Smith . mat - the matrix
1906d3042a70SBarry Smith 
19072ef1f0ffSBarry Smith   Level: developer
19082ef1f0ffSBarry Smith 
190911a5261eSBarry Smith   Note:
191011a5261eSBarry Smith   You can only call this after a call to `MatDensePlaceArray()`
1911d3042a70SBarry Smith 
19121cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDensePlaceArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`
1913d3042a70SBarry Smith @*/
1914d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseResetArray(Mat mat)
1915d71ae5a4SJacob Faibussowitsch {
1916d3042a70SBarry Smith   PetscFunctionBegin;
1917d5ea218eSStefano Zampini   PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
1918cac4c232SBarry Smith   PetscUseMethod(mat, "MatDenseResetArray_C", (Mat), (mat));
19199566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)mat));
19203ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1921d3042a70SBarry Smith }
1922d3042a70SBarry Smith 
1923d5ea218eSStefano Zampini /*@
1924d5ea218eSStefano Zampini   MatDenseReplaceArray - Allows one to replace the array in a dense matrix with an
1925d5ea218eSStefano Zampini   array provided by the user. This is useful to avoid copying an array
1926d5ea218eSStefano Zampini   into a matrix
1927d5ea218eSStefano Zampini 
1928d5ea218eSStefano Zampini   Not Collective
1929d5ea218eSStefano Zampini 
1930d5ea218eSStefano Zampini   Input Parameters:
1931d5ea218eSStefano Zampini + mat   - the matrix
1932d5ea218eSStefano Zampini - array - the array in column major order
1933d5ea218eSStefano Zampini 
19342ef1f0ffSBarry Smith   Level: developer
19352ef1f0ffSBarry Smith 
193611a5261eSBarry Smith   Note:
19374d5500e8SJunchao Zhang   Adding `const` to `array` was an oversight, see notes in `VecPlaceArray()`.
19384d5500e8SJunchao Zhang 
193911a5261eSBarry Smith   The memory passed in MUST be obtained with `PetscMalloc()` and CANNOT be
1940d5ea218eSStefano Zampini   freed by the user. It will be freed when the matrix is destroyed.
1941d5ea218eSStefano Zampini 
19421cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MatDensePlaceArray()`, `MatDenseGetArray()`, `VecReplaceArray()`
1943d5ea218eSStefano Zampini @*/
1944d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseReplaceArray(Mat mat, const PetscScalar *array)
1945d71ae5a4SJacob Faibussowitsch {
1946d5ea218eSStefano Zampini   PetscFunctionBegin;
1947d5ea218eSStefano Zampini   PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
1948cac4c232SBarry Smith   PetscUseMethod(mat, "MatDenseReplaceArray_C", (Mat, const PetscScalar *), (mat, array));
19499566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)mat));
195047d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
1951d5ea218eSStefano Zampini   mat->offloadmask = PETSC_OFFLOAD_CPU;
1952d5ea218eSStefano Zampini #endif
19533ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1954d5ea218eSStefano Zampini }
1955d5ea218eSStefano Zampini 
19565d83a8b1SBarry Smith /*@
195711a5261eSBarry Smith   MatCreateDense - Creates a matrix in `MATDENSE` format.
19588965ea79SLois Curfman McInnes 
1959d083f849SBarry Smith   Collective
1960db81eaa0SLois Curfman McInnes 
19618965ea79SLois Curfman McInnes   Input Parameters:
1962db81eaa0SLois Curfman McInnes + comm - MPI communicator
19632ef1f0ffSBarry Smith . m    - number of local rows (or `PETSC_DECIDE` to have calculated if `M` is given)
19642ef1f0ffSBarry Smith . n    - number of local columns (or `PETSC_DECIDE` to have calculated if `N` is given)
19652ef1f0ffSBarry Smith . M    - number of global rows (or `PETSC_DECIDE` to have calculated if `m` is given)
19662ef1f0ffSBarry Smith . N    - number of global columns (or `PETSC_DECIDE` to have calculated if `n` is given)
19672ef1f0ffSBarry Smith - data - optional location of matrix data.  Set data to `NULL` (`PETSC_NULL_SCALAR` for Fortran users) for PETSc
1968dfc5480cSLois Curfman McInnes    to control all matrix memory allocation.
19698965ea79SLois Curfman McInnes 
19708965ea79SLois Curfman McInnes   Output Parameter:
1971477f1c0bSLois Curfman McInnes . A - the matrix
19728965ea79SLois Curfman McInnes 
19732ef1f0ffSBarry Smith   Level: intermediate
19742ef1f0ffSBarry Smith 
1975b259b22eSLois Curfman McInnes   Notes:
19762ef1f0ffSBarry Smith   The dense format is fully compatible with standard Fortran
197739ddd567SLois Curfman McInnes   storage by columns.
197811a5261eSBarry Smith 
197911a5261eSBarry Smith   Although local portions of the matrix are stored in column-major
1980002bc6daSRichard Tran Mills   order, the matrix is partitioned across MPI ranks by row.
19818965ea79SLois Curfman McInnes 
198218f449edSLois Curfman McInnes   The data input variable is intended primarily for Fortran programmers
198318f449edSLois Curfman McInnes   who wish to allocate their own matrix memory space.  Most users should
19842ef1f0ffSBarry Smith   set `data` to `NULL` (`PETSC_NULL_SCALAR` for Fortran users).
198518f449edSLois Curfman McInnes 
19868965ea79SLois Curfman McInnes   The user MUST specify either the local or global matrix dimensions
19878965ea79SLois Curfman McInnes   (possibly both).
19888965ea79SLois Curfman McInnes 
19891cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()`
19908965ea79SLois Curfman McInnes @*/
1991ce78bad3SBarry Smith PetscErrorCode MatCreateDense(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt M, PetscInt N, PetscScalar data[], Mat *A)
1992d71ae5a4SJacob Faibussowitsch {
19933a40ed3dSBarry Smith   PetscFunctionBegin;
19949566063dSJacob Faibussowitsch   PetscCall(MatCreate(comm, A));
19959566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(*A, m, n, M, N));
19966f08ad05SPierre Jolivet   PetscCall(MatSetType(*A, MATDENSE));
19979566063dSJacob Faibussowitsch   PetscCall(MatSeqDenseSetPreallocation(*A, data));
19986f08ad05SPierre Jolivet   PetscCall(MatMPIDenseSetPreallocation(*A, data));
19993ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
20008965ea79SLois Curfman McInnes }
20018965ea79SLois Curfman McInnes 
2002d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDuplicate_MPIDense(Mat A, MatDuplicateOption cpvalues, Mat *newmat)
2003d71ae5a4SJacob Faibussowitsch {
20048965ea79SLois Curfman McInnes   Mat           mat;
20053501a2bdSLois Curfman McInnes   Mat_MPIDense *a, *oldmat = (Mat_MPIDense *)A->data;
20068965ea79SLois Curfman McInnes 
20073a40ed3dSBarry Smith   PetscFunctionBegin;
2008f4259b30SLisandro Dalcin   *newmat = NULL;
20099566063dSJacob Faibussowitsch   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat));
20109566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(mat, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N));
20119566063dSJacob Faibussowitsch   PetscCall(MatSetType(mat, ((PetscObject)A)->type_name));
2012834f8fabSBarry Smith   a = (Mat_MPIDense *)mat->data;
20135aa7edbeSHong Zhang 
2014d5f3da31SBarry Smith   mat->factortype   = A->factortype;
2015c456f294SBarry Smith   mat->assembled    = PETSC_TRUE;
2016273d9f13SBarry Smith   mat->preallocated = PETSC_TRUE;
20178965ea79SLois Curfman McInnes 
2018e0fa3b82SLois Curfman McInnes   mat->insertmode = NOT_SET_VALUES;
20193782ba37SSatish Balay   a->donotstash   = oldmat->donotstash;
2020e04c1aa4SHong Zhang 
20219566063dSJacob Faibussowitsch   PetscCall(PetscLayoutReference(A->rmap, &mat->rmap));
20229566063dSJacob Faibussowitsch   PetscCall(PetscLayoutReference(A->cmap, &mat->cmap));
20238965ea79SLois Curfman McInnes 
20249566063dSJacob Faibussowitsch   PetscCall(MatDuplicate(oldmat->A, cpvalues, &a->A));
202501b82886SBarry Smith 
20268965ea79SLois Curfman McInnes   *newmat = mat;
20273ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
20288965ea79SLois Curfman McInnes }
20298965ea79SLois Curfman McInnes 
203066976f2fSJacob Faibussowitsch static PetscErrorCode MatLoad_MPIDense(Mat newMat, PetscViewer viewer)
2031d71ae5a4SJacob Faibussowitsch {
203287d5ce66SSatish Balay   PetscBool isbinary;
203387d5ce66SSatish Balay #if defined(PETSC_HAVE_HDF5)
203487d5ce66SSatish Balay   PetscBool ishdf5;
203587d5ce66SSatish Balay #endif
2036eb91f321SVaclav Hapla 
2037eb91f321SVaclav Hapla   PetscFunctionBegin;
2038eb91f321SVaclav Hapla   PetscValidHeaderSpecific(newMat, MAT_CLASSID, 1);
2039eb91f321SVaclav Hapla   PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
2040eb91f321SVaclav Hapla   /* force binary viewer to load .info file if it has not yet done so */
20419566063dSJacob Faibussowitsch   PetscCall(PetscViewerSetUp(viewer));
20429566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
204387d5ce66SSatish Balay #if defined(PETSC_HAVE_HDF5)
20449566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERHDF5, &ishdf5));
204587d5ce66SSatish Balay #endif
2046eb91f321SVaclav Hapla   if (isbinary) {
20479566063dSJacob Faibussowitsch     PetscCall(MatLoad_Dense_Binary(newMat, viewer));
2048eb91f321SVaclav Hapla #if defined(PETSC_HAVE_HDF5)
204987d5ce66SSatish Balay   } else if (ishdf5) {
20509566063dSJacob Faibussowitsch     PetscCall(MatLoad_Dense_HDF5(newMat, viewer));
2051eb91f321SVaclav Hapla #endif
205298921bdaSJacob 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);
20533ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2054eb91f321SVaclav Hapla }
2055eb91f321SVaclav Hapla 
2056d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatEqual_MPIDense(Mat A, Mat B, PetscBool *flag)
2057d71ae5a4SJacob Faibussowitsch {
20586e4ee0c6SHong Zhang   Mat_MPIDense *matB = (Mat_MPIDense *)B->data, *matA = (Mat_MPIDense *)A->data;
20596e4ee0c6SHong Zhang   Mat           a, b;
206090ace30eSBarry Smith 
20616e4ee0c6SHong Zhang   PetscFunctionBegin;
20626e4ee0c6SHong Zhang   a = matA->A;
20636e4ee0c6SHong Zhang   b = matB->A;
20642cf15c64SPierre Jolivet   PetscCall(MatEqual(a, b, flag));
2065462c564dSBarry Smith   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, flag, 1, MPIU_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
20663ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
20676e4ee0c6SHong Zhang }
206890ace30eSBarry Smith 
206966976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MatTransMatMult_MPIDense_MPIDense(void *data)
2070d71ae5a4SJacob Faibussowitsch {
20716718818eSStefano Zampini   Mat_TransMatMultDense *atb = (Mat_TransMatMultDense *)data;
2072baa3c1c6SHong Zhang 
2073baa3c1c6SHong Zhang   PetscFunctionBegin;
20749566063dSJacob Faibussowitsch   PetscCall(PetscFree2(atb->sendbuf, atb->recvcounts));
20759566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&atb->atb));
20769566063dSJacob Faibussowitsch   PetscCall(PetscFree(atb));
20773ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2078baa3c1c6SHong Zhang }
2079baa3c1c6SHong Zhang 
208066976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MatMatTransMult_MPIDense_MPIDense(void *data)
2081d71ae5a4SJacob Faibussowitsch {
20826718818eSStefano Zampini   Mat_MatTransMultDense *abt = (Mat_MatTransMultDense *)data;
2083cc48ffa7SToby Isaac 
2084cc48ffa7SToby Isaac   PetscFunctionBegin;
20859566063dSJacob Faibussowitsch   PetscCall(PetscFree2(abt->buf[0], abt->buf[1]));
20869566063dSJacob Faibussowitsch   PetscCall(PetscFree2(abt->recvcounts, abt->recvdispls));
20879566063dSJacob Faibussowitsch   PetscCall(PetscFree(abt));
20883ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2089cc48ffa7SToby Isaac }
2090cc48ffa7SToby Isaac 
2091d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2092d71ae5a4SJacob Faibussowitsch {
2093baa3c1c6SHong Zhang   Mat_MPIDense          *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
20946718818eSStefano Zampini   Mat_TransMatMultDense *atb;
2095cb20be35SHong Zhang   MPI_Comm               comm;
20966718818eSStefano Zampini   PetscMPIInt            size, *recvcounts;
20976718818eSStefano Zampini   PetscScalar           *carray, *sendbuf;
2098637a0070SStefano Zampini   const PetscScalar     *atbarray;
2099a1f9a5e6SStefano Zampini   PetscInt               i, cN = C->cmap->N, proc, k, j, lda;
2100e68c0b26SHong Zhang   const PetscInt        *ranges;
2101cb20be35SHong Zhang 
2102cb20be35SHong Zhang   PetscFunctionBegin;
21036718818eSStefano Zampini   MatCheckProduct(C, 3);
21045f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
21056718818eSStefano Zampini   atb        = (Mat_TransMatMultDense *)C->product->data;
21066718818eSStefano Zampini   recvcounts = atb->recvcounts;
21076718818eSStefano Zampini   sendbuf    = atb->sendbuf;
21086718818eSStefano Zampini 
21099566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
21109566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
2111e68c0b26SHong Zhang 
2112c5ef1628SHong Zhang   /* compute atbarray = aseq^T * bseq */
2113fb842aefSJose E. Roman   PetscCall(MatTransposeMatMult(a->A, b->A, atb->atb ? MAT_REUSE_MATRIX : MAT_INITIAL_MATRIX, PETSC_DETERMINE, &atb->atb));
2114cb20be35SHong Zhang 
21159566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRanges(C, &ranges));
2116cb20be35SHong Zhang 
2117d016bddeSToby Isaac   if (ranges[1] == C->rmap->N) {
2118d016bddeSToby Isaac     /* all of the values are being reduced to rank 0: optimize this case to use MPI_Reduce and GPU aware MPI if available */
2119d016bddeSToby Isaac     PetscInt           atb_lda, c_lda;
2120d016bddeSToby Isaac     Mat                atb_local = atb->atb;
2121d016bddeSToby Isaac     Mat                atb_alloc = NULL;
2122d016bddeSToby Isaac     Mat                c_local   = c->A;
2123d016bddeSToby Isaac     Mat                c_alloc   = NULL;
2124d016bddeSToby Isaac     PetscMemType       atb_memtype, c_memtype;
2125d016bddeSToby Isaac     const PetscScalar *atb_array = NULL;
2126d016bddeSToby Isaac     MPI_Datatype       vector_type;
2127d016bddeSToby Isaac     PetscScalar       *c_array = NULL;
2128d016bddeSToby Isaac     PetscMPIInt        rank;
2129d016bddeSToby Isaac 
2130d016bddeSToby Isaac     PetscCallMPI(MPI_Comm_rank(comm, &rank));
2131d016bddeSToby Isaac 
2132d016bddeSToby Isaac     PetscCall(MatDenseGetLDA(atb_local, &atb_lda));
2133d016bddeSToby Isaac     if (atb_lda != C->rmap->N) {
2134d016bddeSToby Isaac       // copy atb to a matrix that will have lda == the number of rows
2135d016bddeSToby Isaac       PetscCall(MatDuplicate(atb_local, MAT_DO_NOT_COPY_VALUES, &atb_alloc));
2136d016bddeSToby Isaac       PetscCall(MatCopy(atb_local, atb_alloc, DIFFERENT_NONZERO_PATTERN));
2137d016bddeSToby Isaac       atb_local = atb_alloc;
2138d016bddeSToby Isaac     }
2139d016bddeSToby Isaac 
2140d016bddeSToby Isaac     if (rank == 0) {
2141d016bddeSToby Isaac       PetscCall(MatDenseGetLDA(c_local, &c_lda));
2142d016bddeSToby Isaac       if (c_lda != C->rmap->N) {
2143d016bddeSToby Isaac         // copy c to a matrix that will have lda == the number of rows
2144d016bddeSToby Isaac         PetscCall(MatDuplicate(c_local, MAT_DO_NOT_COPY_VALUES, &c_alloc));
2145d016bddeSToby Isaac         c_local = c_alloc;
2146d016bddeSToby Isaac       }
2147d016bddeSToby Isaac       PetscCall(MatZeroEntries(c_local));
2148d016bddeSToby Isaac     }
2149d016bddeSToby Isaac     /* atb_local and c_local have nrows = lda = A->cmap->N and ncols =
2150d016bddeSToby Isaac      * B->cmap->N: use the a->Mvctx to use the best reduction method */
2151d016bddeSToby Isaac     if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
2152d016bddeSToby Isaac     vector_type = MPIU_SCALAR;
2153d016bddeSToby Isaac     if (B->cmap->N > 1) {
2154*b682a77bSToby Isaac       PetscMPIInt mpi_N;
2155*b682a77bSToby Isaac 
2156*b682a77bSToby Isaac       PetscCall(PetscMPIIntCast(B->cmap->N, &mpi_N));
2157*b682a77bSToby Isaac       PetscCallMPI(MPI_Type_contiguous(mpi_N, MPIU_SCALAR, &vector_type));
2158d016bddeSToby Isaac       PetscCallMPI(MPI_Type_commit(&vector_type));
2159d016bddeSToby Isaac     }
2160d016bddeSToby Isaac     PetscCall(MatDenseGetArrayReadAndMemType(atb_local, &atb_array, &atb_memtype));
2161d016bddeSToby Isaac     PetscCall(MatDenseGetArrayWriteAndMemType(c_local, &c_array, &c_memtype));
2162d016bddeSToby Isaac     PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, vector_type, atb_memtype, atb_array, c_memtype, c_array, MPIU_SUM));
2163d016bddeSToby Isaac     PetscCall(PetscSFReduceEnd(a->Mvctx, vector_type, atb_array, c_array, MPIU_SUM));
2164d016bddeSToby Isaac     PetscCall(MatDenseRestoreArrayWriteAndMemType(c_local, &c_array));
2165d016bddeSToby Isaac     PetscCall(MatDenseRestoreArrayReadAndMemType(atb_local, &atb_array));
2166d016bddeSToby Isaac     if (rank == 0 && c_local != c->A) PetscCall(MatCopy(c_local, c->A, DIFFERENT_NONZERO_PATTERN));
2167d016bddeSToby Isaac     if (B->cmap->N > 1) PetscCallMPI(MPI_Type_free(&vector_type));
2168d016bddeSToby Isaac     PetscCall(MatDestroy(&atb_alloc));
2169d016bddeSToby Isaac     PetscCall(MatDestroy(&c_alloc));
2170d016bddeSToby Isaac     PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2171d016bddeSToby Isaac     PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2172d016bddeSToby Isaac     PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2173d016bddeSToby Isaac     PetscFunctionReturn(PETSC_SUCCESS);
2174d016bddeSToby Isaac   }
2175d016bddeSToby Isaac 
2176660d5466SHong Zhang   /* arrange atbarray into sendbuf */
21779566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(atb->atb, &atbarray));
21789566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(atb->atb, &lda));
2179637a0070SStefano Zampini   for (proc = 0, k = 0; proc < size; proc++) {
2180baa3c1c6SHong Zhang     for (j = 0; j < cN; j++) {
2181a1f9a5e6SStefano Zampini       for (i = ranges[proc]; i < ranges[proc + 1]; i++) sendbuf[k++] = atbarray[i + j * lda];
2182cb20be35SHong Zhang     }
2183cb20be35SHong Zhang   }
21849566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(atb->atb, &atbarray));
2185637a0070SStefano Zampini 
2186c5ef1628SHong Zhang   /* sum all atbarray to local values of C */
21879566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(c->A, &carray));
21889566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Reduce_scatter(sendbuf, carray, recvcounts, MPIU_SCALAR, MPIU_SUM, comm));
21899566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(c->A, &carray));
219067af85e8SPierre Jolivet   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
21919566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
21929566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
21933ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2194cb20be35SHong Zhang }
2195cb20be35SHong Zhang 
2196d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2197d71ae5a4SJacob Faibussowitsch {
2198cb20be35SHong Zhang   MPI_Comm               comm;
2199baa3c1c6SHong Zhang   PetscMPIInt            size;
2200660d5466SHong Zhang   PetscInt               cm = A->cmap->n, cM, cN = B->cmap->N;
2201baa3c1c6SHong Zhang   Mat_TransMatMultDense *atb;
220247d993e7Ssuyashtn   PetscBool              cisdense = PETSC_FALSE;
22030acaeb71SStefano Zampini   const PetscInt        *ranges;
2204cb20be35SHong Zhang 
2205cb20be35SHong Zhang   PetscFunctionBegin;
22068f83a4d9SJacob Faibussowitsch   MatCheckProduct(C, 4);
22075f80ce2aSJacob Faibussowitsch   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
22089566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
2209cb20be35SHong Zhang   if (A->rmap->rstart != B->rmap->rstart || A->rmap->rend != B->rmap->rend) {
221098921bdaSJacob Faibussowitsch     SETERRQ(comm, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != B (%" PetscInt_FMT ",%" PetscInt_FMT ")", A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
2211cb20be35SHong Zhang   }
2212cb20be35SHong Zhang 
22134222ddf1SHong Zhang   /* create matrix product C */
22149566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(C, cm, B->cmap->n, A->cmap->N, B->cmap->N));
221547d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA)
22169566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSECUDA, ""));
221747d993e7Ssuyashtn #endif
221847d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
221947d993e7Ssuyashtn   PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSEHIP, ""));
222047d993e7Ssuyashtn #endif
222148a46eb9SPierre Jolivet   if (!cisdense) PetscCall(MatSetType(C, ((PetscObject)A)->type_name));
22229566063dSJacob Faibussowitsch   PetscCall(MatSetUp(C));
2223baa3c1c6SHong Zhang 
22244222ddf1SHong Zhang   /* create data structure for reuse C */
22259566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
22269566063dSJacob Faibussowitsch   PetscCall(PetscNew(&atb));
22274222ddf1SHong Zhang   cM = C->rmap->N;
22289566063dSJacob Faibussowitsch   PetscCall(PetscMalloc2(cM * cN, &atb->sendbuf, size, &atb->recvcounts));
22299566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRanges(C, &ranges));
22306497c311SBarry Smith   for (PetscMPIInt i = 0; i < size; i++) PetscCall(PetscMPIIntCast((ranges[i + 1] - ranges[i]) * cN, &atb->recvcounts[i]));
22316718818eSStefano Zampini   C->product->data    = atb;
22326718818eSStefano Zampini   C->product->destroy = MatDestroy_MatTransMatMult_MPIDense_MPIDense;
22333ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2234cb20be35SHong Zhang }
2235cb20be35SHong Zhang 
2236d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2237d71ae5a4SJacob Faibussowitsch {
2238cc48ffa7SToby Isaac   MPI_Comm               comm;
2239cc48ffa7SToby Isaac   PetscMPIInt            i, size;
2240cc48ffa7SToby Isaac   PetscInt               maxRows, bufsiz;
2241cc48ffa7SToby Isaac   PetscMPIInt            tag;
22424222ddf1SHong Zhang   PetscInt               alg;
2243cc48ffa7SToby Isaac   Mat_MatTransMultDense *abt;
22444222ddf1SHong Zhang   Mat_Product           *product = C->product;
22454222ddf1SHong Zhang   PetscBool              flg;
2246cc48ffa7SToby Isaac 
2247cc48ffa7SToby Isaac   PetscFunctionBegin;
22486718818eSStefano Zampini   MatCheckProduct(C, 4);
22495f80ce2aSJacob Faibussowitsch   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
22504222ddf1SHong Zhang   /* check local size of A and B */
22515f80ce2aSJacob 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);
2252cc48ffa7SToby Isaac 
22539566063dSJacob Faibussowitsch   PetscCall(PetscStrcmp(product->alg, "allgatherv", &flg));
2254637a0070SStefano Zampini   alg = flg ? 0 : 1;
2255cc48ffa7SToby Isaac 
22564222ddf1SHong Zhang   /* setup matrix product C */
22579566063dSJacob Faibussowitsch   PetscCall(MatSetSizes(C, A->rmap->n, B->rmap->n, A->rmap->N, B->rmap->N));
22589566063dSJacob Faibussowitsch   PetscCall(MatSetType(C, MATMPIDENSE));
22599566063dSJacob Faibussowitsch   PetscCall(MatSetUp(C));
22609566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag));
22614222ddf1SHong Zhang 
22624222ddf1SHong Zhang   /* create data structure for reuse C */
22639566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
22649566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
22659566063dSJacob Faibussowitsch   PetscCall(PetscNew(&abt));
2266cc48ffa7SToby Isaac   abt->tag = tag;
2267faa55883SToby Isaac   abt->alg = alg;
2268faa55883SToby Isaac   switch (alg) {
22694222ddf1SHong Zhang   case 1: /* alg: "cyclic" */
227057508eceSPierre Jolivet     for (maxRows = 0, i = 0; i < size; i++) maxRows = PetscMax(maxRows, B->rmap->range[i + 1] - B->rmap->range[i]);
2271cc48ffa7SToby Isaac     bufsiz = A->cmap->N * maxRows;
2272f4f49eeaSPierre Jolivet     PetscCall(PetscMalloc2(bufsiz, &abt->buf[0], bufsiz, &abt->buf[1]));
2273faa55883SToby Isaac     break;
22744222ddf1SHong Zhang   default: /* alg: "allgatherv" */
2275f4f49eeaSPierre Jolivet     PetscCall(PetscMalloc2(B->rmap->n * B->cmap->N, &abt->buf[0], B->rmap->N * B->cmap->N, &abt->buf[1]));
2276f4f49eeaSPierre Jolivet     PetscCall(PetscMalloc2(size, &abt->recvcounts, size + 1, &abt->recvdispls));
22776497c311SBarry Smith     for (i = 0; i <= size; i++) PetscCall(PetscMPIIntCast(B->rmap->range[i] * A->cmap->N, &abt->recvdispls[i]));
22786497c311SBarry Smith     for (i = 0; i < size; i++) PetscCall(PetscMPIIntCast(abt->recvdispls[i + 1] - abt->recvdispls[i], &abt->recvcounts[i]));
2279faa55883SToby Isaac     break;
2280faa55883SToby Isaac   }
2281cc48ffa7SToby Isaac 
22826718818eSStefano Zampini   C->product->data                = abt;
22836718818eSStefano Zampini   C->product->destroy             = MatDestroy_MatMatTransMult_MPIDense_MPIDense;
22844222ddf1SHong Zhang   C->ops->mattransposemultnumeric = MatMatTransposeMultNumeric_MPIDense_MPIDense;
22853ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2286cc48ffa7SToby Isaac }
2287cc48ffa7SToby Isaac 
2288d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(Mat A, Mat B, Mat C)
2289d71ae5a4SJacob Faibussowitsch {
2290cc48ffa7SToby Isaac   Mat_MPIDense          *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
22916718818eSStefano Zampini   Mat_MatTransMultDense *abt;
2292cc48ffa7SToby Isaac   MPI_Comm               comm;
22936497c311SBarry Smith   PetscMPIInt            rank, size, sendto, recvfrom, recvisfrom;
2294f4259b30SLisandro Dalcin   PetscScalar           *sendbuf, *recvbuf = NULL, *cv;
22956497c311SBarry Smith   PetscInt               i, cK             = A->cmap->N, sendsiz, recvsiz, k, j, bn;
2296cc48ffa7SToby Isaac   PetscScalar            _DOne = 1.0, _DZero = 0.0;
2297637a0070SStefano Zampini   const PetscScalar     *av, *bv;
229823fff9afSBarry Smith   PetscBLASInt           cm, cn, ck, alda, blda = 0, clda;
2299cc48ffa7SToby Isaac   MPI_Request            reqs[2];
2300cc48ffa7SToby Isaac   const PetscInt        *ranges;
2301cc48ffa7SToby Isaac 
2302cc48ffa7SToby Isaac   PetscFunctionBegin;
23036718818eSStefano Zampini   MatCheckProduct(C, 3);
23045f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
23056718818eSStefano Zampini   abt = (Mat_MatTransMultDense *)C->product->data;
23069566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
23079566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(comm, &rank));
23089566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
23099566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, &av));
23109566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(b->A, &bv));
23119566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(c->A, &cv));
23129566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &i));
23139566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &alda));
23149566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(b->A, &i));
23159566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &blda));
23169566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(c->A, &i));
23179566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &clda));
23189566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRanges(B, &ranges));
2319cc48ffa7SToby Isaac   bn = B->rmap->n;
2320637a0070SStefano Zampini   if (blda == bn) {
2321637a0070SStefano Zampini     sendbuf = (PetscScalar *)bv;
2322cc48ffa7SToby Isaac   } else {
2323cc48ffa7SToby Isaac     sendbuf = abt->buf[0];
2324cc48ffa7SToby Isaac     for (k = 0, i = 0; i < cK; i++) {
2325ad540459SPierre Jolivet       for (j = 0; j < bn; j++, k++) sendbuf[k] = bv[i * blda + j];
2326cc48ffa7SToby Isaac     }
2327cc48ffa7SToby Isaac   }
2328cc48ffa7SToby Isaac   if (size > 1) {
2329cc48ffa7SToby Isaac     sendto   = (rank + size - 1) % size;
2330cc48ffa7SToby Isaac     recvfrom = (rank + size + 1) % size;
2331cc48ffa7SToby Isaac   } else {
2332cc48ffa7SToby Isaac     sendto = recvfrom = 0;
2333cc48ffa7SToby Isaac   }
23349566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(cK, &ck));
23359566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm));
2336cc48ffa7SToby Isaac   recvisfrom = rank;
2337cc48ffa7SToby Isaac   for (i = 0; i < size; i++) {
2338cc48ffa7SToby Isaac     /* we have finished receiving in sending, bufs can be read/modified */
23396497c311SBarry Smith     PetscMPIInt nextrecvisfrom = (recvisfrom + 1) % size; /* which process the next recvbuf will originate on */
2340cc48ffa7SToby Isaac     PetscInt    nextbn         = ranges[nextrecvisfrom + 1] - ranges[nextrecvisfrom];
2341cc48ffa7SToby Isaac 
2342cc48ffa7SToby Isaac     if (nextrecvisfrom != rank) {
2343cc48ffa7SToby Isaac       /* start the cyclic sends from sendbuf, to recvbuf (which will switch to sendbuf) */
2344cc48ffa7SToby Isaac       sendsiz = cK * bn;
2345cc48ffa7SToby Isaac       recvsiz = cK * nextbn;
2346cc48ffa7SToby Isaac       recvbuf = (i & 1) ? abt->buf[0] : abt->buf[1];
23476497c311SBarry Smith       PetscCallMPI(MPIU_Isend(sendbuf, sendsiz, MPIU_SCALAR, sendto, abt->tag, comm, &reqs[0]));
23486497c311SBarry Smith       PetscCallMPI(MPIU_Irecv(recvbuf, recvsiz, MPIU_SCALAR, recvfrom, abt->tag, comm, &reqs[1]));
2349cc48ffa7SToby Isaac     }
2350cc48ffa7SToby Isaac 
2351cc48ffa7SToby Isaac     /* local aseq * sendbuf^T */
23529566063dSJacob Faibussowitsch     PetscCall(PetscBLASIntCast(ranges[recvisfrom + 1] - ranges[recvisfrom], &cn));
2353792fecdfSBarry Smith     if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "T", &cm, &cn, &ck, &_DOne, av, &alda, sendbuf, &cn, &_DZero, cv + clda * ranges[recvisfrom], &clda));
2354cc48ffa7SToby Isaac 
2355cc48ffa7SToby Isaac     if (nextrecvisfrom != rank) {
2356cc48ffa7SToby Isaac       /* wait for the sends and receives to complete, swap sendbuf and recvbuf */
23579566063dSJacob Faibussowitsch       PetscCallMPI(MPI_Waitall(2, reqs, MPI_STATUSES_IGNORE));
2358cc48ffa7SToby Isaac     }
2359cc48ffa7SToby Isaac     bn         = nextbn;
2360cc48ffa7SToby Isaac     recvisfrom = nextrecvisfrom;
2361cc48ffa7SToby Isaac     sendbuf    = recvbuf;
2362cc48ffa7SToby Isaac   }
23639566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, &av));
23649566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
23659566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(c->A, &cv));
236667af85e8SPierre Jolivet   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
23679566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
23689566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
23693ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2370cc48ffa7SToby Isaac }
2371cc48ffa7SToby Isaac 
2372d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(Mat A, Mat B, Mat C)
2373d71ae5a4SJacob Faibussowitsch {
2374faa55883SToby Isaac   Mat_MPIDense          *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
23756718818eSStefano Zampini   Mat_MatTransMultDense *abt;
2376faa55883SToby Isaac   MPI_Comm               comm;
23776497c311SBarry Smith   PetscMPIInt            size, ibn;
2378637a0070SStefano Zampini   PetscScalar           *cv, *sendbuf, *recvbuf;
2379637a0070SStefano Zampini   const PetscScalar     *av, *bv;
2380637a0070SStefano Zampini   PetscInt               blda, i, cK = A->cmap->N, k, j, bn;
2381faa55883SToby Isaac   PetscScalar            _DOne = 1.0, _DZero = 0.0;
2382637a0070SStefano Zampini   PetscBLASInt           cm, cn, ck, alda, clda;
2383faa55883SToby Isaac 
2384faa55883SToby Isaac   PetscFunctionBegin;
23856718818eSStefano Zampini   MatCheckProduct(C, 3);
23865f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
23876718818eSStefano Zampini   abt = (Mat_MatTransMultDense *)C->product->data;
23889566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
23899566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_size(comm, &size));
23909566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(a->A, &av));
23919566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayRead(b->A, &bv));
23929566063dSJacob Faibussowitsch   PetscCall(MatDenseGetArrayWrite(c->A, &cv));
23939566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(a->A, &i));
23949566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &alda));
23959566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(b->A, &blda));
23969566063dSJacob Faibussowitsch   PetscCall(MatDenseGetLDA(c->A, &i));
23979566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(i, &clda));
2398faa55883SToby Isaac   /* copy transpose of B into buf[0] */
2399faa55883SToby Isaac   bn      = B->rmap->n;
2400faa55883SToby Isaac   sendbuf = abt->buf[0];
2401faa55883SToby Isaac   recvbuf = abt->buf[1];
2402faa55883SToby Isaac   for (k = 0, j = 0; j < bn; j++) {
2403ad540459SPierre Jolivet     for (i = 0; i < cK; i++, k++) sendbuf[k] = bv[i * blda + j];
2404faa55883SToby Isaac   }
24059566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
24066497c311SBarry Smith   PetscCall(PetscMPIIntCast(bn * cK, &ibn));
24076497c311SBarry Smith   PetscCallMPI(MPI_Allgatherv(sendbuf, ibn, MPIU_SCALAR, recvbuf, abt->recvcounts, abt->recvdispls, MPIU_SCALAR, comm));
24089566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(cK, &ck));
24099566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm));
24109566063dSJacob Faibussowitsch   PetscCall(PetscBLASIntCast(c->A->cmap->n, &cn));
2411792fecdfSBarry Smith   if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "N", &cm, &cn, &ck, &_DOne, av, &alda, recvbuf, &ck, &_DZero, cv, &clda));
24129566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(a->A, &av));
24139566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
24149566063dSJacob Faibussowitsch   PetscCall(MatDenseRestoreArrayWrite(c->A, &cv));
241567af85e8SPierre Jolivet   PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
24169566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
24179566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
24183ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2419faa55883SToby Isaac }
2420faa55883SToby Isaac 
2421d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2422d71ae5a4SJacob Faibussowitsch {
24236718818eSStefano Zampini   Mat_MatTransMultDense *abt;
2424faa55883SToby Isaac 
2425faa55883SToby Isaac   PetscFunctionBegin;
24266718818eSStefano Zampini   MatCheckProduct(C, 3);
24275f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
24286718818eSStefano Zampini   abt = (Mat_MatTransMultDense *)C->product->data;
2429faa55883SToby Isaac   switch (abt->alg) {
2430d71ae5a4SJacob Faibussowitsch   case 1:
2431d71ae5a4SJacob Faibussowitsch     PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(A, B, C));
2432d71ae5a4SJacob Faibussowitsch     break;
2433d71ae5a4SJacob Faibussowitsch   default:
2434d71ae5a4SJacob Faibussowitsch     PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(A, B, C));
2435d71ae5a4SJacob Faibussowitsch     break;
2436faa55883SToby Isaac   }
24373ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2438faa55883SToby Isaac }
2439faa55883SToby Isaac 
244066976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MatMatMult_MPIDense_MPIDense(void *data)
2441d71ae5a4SJacob Faibussowitsch {
24426718818eSStefano Zampini   Mat_MatMultDense *ab = (Mat_MatMultDense *)data;
2443320f2790SHong Zhang 
2444320f2790SHong Zhang   PetscFunctionBegin;
24459566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&ab->Ce));
24469566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&ab->Ae));
24479566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&ab->Be));
24489566063dSJacob Faibussowitsch   PetscCall(PetscFree(ab));
24493ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2450320f2790SHong Zhang }
2451320f2790SHong Zhang 
24525d8c7819SPierre Jolivet static PetscErrorCode MatMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2453d71ae5a4SJacob Faibussowitsch {
24546718818eSStefano Zampini   Mat_MatMultDense *ab;
24555d8c7819SPierre Jolivet   Mat_MPIDense     *mdn = (Mat_MPIDense *)A->data;
2456d016bddeSToby Isaac   Mat_MPIDense     *b   = (Mat_MPIDense *)B->data;
2457320f2790SHong Zhang 
2458320f2790SHong Zhang   PetscFunctionBegin;
24596718818eSStefano Zampini   MatCheckProduct(C, 3);
24605f80ce2aSJacob Faibussowitsch   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Missing product data");
24616718818eSStefano Zampini   ab = (Mat_MatMultDense *)C->product->data;
24625d8c7819SPierre Jolivet   if (ab->Ae && ab->Ce) {
24635d8c7819SPierre Jolivet #if PetscDefined(HAVE_ELEMENTAL)
24649566063dSJacob Faibussowitsch     PetscCall(MatConvert_MPIDense_Elemental(A, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Ae));
24659566063dSJacob Faibussowitsch     PetscCall(MatConvert_MPIDense_Elemental(B, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Be));
24669566063dSJacob Faibussowitsch     PetscCall(MatMatMultNumeric_Elemental(ab->Ae, ab->Be, ab->Ce));
24679566063dSJacob Faibussowitsch     PetscCall(MatConvert(ab->Ce, MATMPIDENSE, MAT_REUSE_MATRIX, &C));
24685d8c7819SPierre Jolivet #else
24695d8c7819SPierre Jolivet     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined");
24705d8c7819SPierre Jolivet #endif
24715d8c7819SPierre Jolivet   } else {
2472d016bddeSToby Isaac     MPI_Comm           comm;
24735d8c7819SPierre Jolivet     const PetscScalar *read;
24745d8c7819SPierre Jolivet     PetscScalar       *write;
24755d8c7819SPierre Jolivet     PetscInt           lda;
2476d016bddeSToby Isaac     const PetscInt    *ranges;
2477d016bddeSToby Isaac     PetscMPIInt        size;
24785d8c7819SPierre Jolivet 
2479d016bddeSToby Isaac     if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A)); /* cannot be done during the symbolic phase because of possible calls to MatProductReplaceMats() */
2480d016bddeSToby Isaac     comm = PetscObjectComm((PetscObject)B);
2481d016bddeSToby Isaac     PetscCallMPI(MPI_Comm_size(comm, &size));
2482d016bddeSToby Isaac     PetscCall(PetscLayoutGetRanges(B->rmap, &ranges));
2483d016bddeSToby Isaac     if (ranges[1] == ranges[size]) {
2484d016bddeSToby Isaac       // optimize for the case where the B matrix is broadcast from rank 0
2485d016bddeSToby Isaac       PetscInt           b_lda, be_lda;
2486d016bddeSToby Isaac       Mat                b_local  = b->A;
2487d016bddeSToby Isaac       Mat                b_alloc  = NULL;
2488d016bddeSToby Isaac       Mat                be_local = ab->Be;
2489d016bddeSToby Isaac       Mat                be_alloc = NULL;
2490d016bddeSToby Isaac       PetscMemType       b_memtype, be_memtype;
2491d016bddeSToby Isaac       const PetscScalar *b_array = NULL;
2492d016bddeSToby Isaac       MPI_Datatype       vector_type;
2493d016bddeSToby Isaac       PetscScalar       *be_array = NULL;
2494d016bddeSToby Isaac       PetscMPIInt        rank;
2495d016bddeSToby Isaac 
2496d016bddeSToby Isaac       PetscCallMPI(MPI_Comm_rank(comm, &rank));
2497d016bddeSToby Isaac       PetscCall(MatDenseGetLDA(be_local, &be_lda));
2498d016bddeSToby Isaac       if (be_lda != B->rmap->N) {
2499d016bddeSToby Isaac         PetscCall(MatDuplicate(be_local, MAT_DO_NOT_COPY_VALUES, &be_alloc));
2500d016bddeSToby Isaac         be_local = be_alloc;
2501d016bddeSToby Isaac       }
2502d016bddeSToby Isaac 
2503d016bddeSToby Isaac       if (rank == 0) {
2504d016bddeSToby Isaac         PetscCall(MatDenseGetLDA(b_local, &b_lda));
2505d016bddeSToby Isaac         if (b_lda != B->rmap->N) {
2506d016bddeSToby Isaac           PetscCall(MatDuplicate(b_local, MAT_DO_NOT_COPY_VALUES, &b_alloc));
2507d016bddeSToby Isaac           PetscCall(MatCopy(b_local, b_alloc, DIFFERENT_NONZERO_PATTERN));
2508d016bddeSToby Isaac           b_local = b_alloc;
2509d016bddeSToby Isaac         }
2510d016bddeSToby Isaac       }
2511*b682a77bSToby Isaac       vector_type = MPIU_SCALAR;
2512*b682a77bSToby Isaac       if (B->cmap->N > 1) {
2513*b682a77bSToby Isaac         PetscMPIInt mpi_N;
2514*b682a77bSToby Isaac 
2515*b682a77bSToby Isaac         PetscCall(PetscMPIIntCast(B->cmap->N, &mpi_N));
2516*b682a77bSToby Isaac         PetscCallMPI(MPI_Type_contiguous(mpi_N, MPIU_SCALAR, &vector_type));
2517d016bddeSToby Isaac         PetscCallMPI(MPI_Type_commit(&vector_type));
2518*b682a77bSToby Isaac       }
2519d016bddeSToby Isaac       PetscCall(MatDenseGetArrayReadAndMemType(b_local, &b_array, &b_memtype));
2520d016bddeSToby Isaac       PetscCall(MatDenseGetArrayWriteAndMemType(be_local, &be_array, &be_memtype));
2521d016bddeSToby Isaac       PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, vector_type, b_memtype, b_array, be_memtype, be_array, MPI_REPLACE));
2522d016bddeSToby Isaac       PetscCall(PetscSFBcastEnd(mdn->Mvctx, vector_type, b_array, be_array, MPI_REPLACE));
2523d016bddeSToby Isaac       PetscCall(MatDenseRestoreArrayWriteAndMemType(be_local, &be_array));
2524d016bddeSToby Isaac       PetscCall(MatDenseRestoreArrayReadAndMemType(b_local, &b_array));
2525d016bddeSToby Isaac       if (be_local != ab->Be) PetscCall(MatCopy(be_local, ab->Be, DIFFERENT_NONZERO_PATTERN));
2526*b682a77bSToby Isaac       if (B->cmap->N > 1) PetscCallMPI(MPI_Type_free(&vector_type));
2527d016bddeSToby Isaac       PetscCall(MatDestroy(&be_alloc));
2528d016bddeSToby Isaac       PetscCall(MatDestroy(&b_alloc));
2529d016bddeSToby Isaac     } else {
25305d8c7819SPierre Jolivet       PetscCall(MatDenseGetLDA(B, &lda));
25315d8c7819SPierre Jolivet       PetscCall(MatDenseGetArrayRead(B, &read));
25325d8c7819SPierre Jolivet       PetscCall(MatDenseGetArrayWrite(ab->Be, &write));
25335d8c7819SPierre Jolivet       for (PetscInt i = 0; i < C->cmap->N; ++i) {
25345d8c7819SPierre Jolivet         PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE));
25355d8c7819SPierre Jolivet         PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE));
25365d8c7819SPierre Jolivet       }
25375d8c7819SPierre Jolivet       PetscCall(MatDenseRestoreArrayWrite(ab->Be, &write));
25385d8c7819SPierre Jolivet       PetscCall(MatDenseRestoreArrayRead(B, &read));
2539d016bddeSToby Isaac     }
2540f4f49eeaSPierre Jolivet     PetscCall(MatMatMultNumeric_SeqDense_SeqDense(((Mat_MPIDense *)A->data)->A, ab->Be, ((Mat_MPIDense *)C->data)->A));
25415d8c7819SPierre Jolivet   }
25423ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2543320f2790SHong Zhang }
2544320f2790SHong Zhang 
2545d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2546d71ae5a4SJacob Faibussowitsch {
25475d8c7819SPierre Jolivet   Mat_Product      *product = C->product;
25485d8c7819SPierre Jolivet   PetscInt          alg;
2549320f2790SHong Zhang   Mat_MatMultDense *ab;
25505d8c7819SPierre Jolivet   PetscBool         flg;
2551320f2790SHong Zhang 
2552320f2790SHong Zhang   PetscFunctionBegin;
25536718818eSStefano Zampini   MatCheckProduct(C, 4);
25545f80ce2aSJacob Faibussowitsch   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
25554222ddf1SHong Zhang   /* check local size of A and B */
25565d8c7819SPierre 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 ")",
25575d8c7819SPierre Jolivet              A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
2558320f2790SHong Zhang 
25595d8c7819SPierre Jolivet   PetscCall(PetscStrcmp(product->alg, "petsc", &flg));
25605d8c7819SPierre Jolivet   alg = flg ? 0 : 1;
25614222ddf1SHong Zhang 
25624222ddf1SHong Zhang   /* setup C */
25635d8c7819SPierre Jolivet   PetscCall(MatSetSizes(C, A->rmap->n, B->cmap->n, A->rmap->N, B->cmap->N));
25645d8c7819SPierre Jolivet   PetscCall(MatSetType(C, MATMPIDENSE));
25659566063dSJacob Faibussowitsch   PetscCall(MatSetUp(C));
2566320f2790SHong Zhang 
2567320f2790SHong Zhang   /* create data structure for reuse Cdense */
25689566063dSJacob Faibussowitsch   PetscCall(PetscNew(&ab));
25695d8c7819SPierre Jolivet 
25705d8c7819SPierre Jolivet   switch (alg) {
25715d8c7819SPierre Jolivet   case 1: /* alg: "elemental" */
25725d8c7819SPierre Jolivet #if PetscDefined(HAVE_ELEMENTAL)
25735d8c7819SPierre Jolivet     /* create elemental matrices Ae and Be */
25745d8c7819SPierre Jolivet     PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &ab->Ae));
25755d8c7819SPierre Jolivet     PetscCall(MatSetSizes(ab->Ae, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N));
25765d8c7819SPierre Jolivet     PetscCall(MatSetType(ab->Ae, MATELEMENTAL));
25775d8c7819SPierre Jolivet     PetscCall(MatSetUp(ab->Ae));
25785d8c7819SPierre Jolivet     PetscCall(MatSetOption(ab->Ae, MAT_ROW_ORIENTED, PETSC_FALSE));
25795d8c7819SPierre Jolivet 
25805d8c7819SPierre Jolivet     PetscCall(MatCreate(PetscObjectComm((PetscObject)B), &ab->Be));
25815d8c7819SPierre Jolivet     PetscCall(MatSetSizes(ab->Be, PETSC_DECIDE, PETSC_DECIDE, B->rmap->N, B->cmap->N));
25825d8c7819SPierre Jolivet     PetscCall(MatSetType(ab->Be, MATELEMENTAL));
25835d8c7819SPierre Jolivet     PetscCall(MatSetUp(ab->Be));
25845d8c7819SPierre Jolivet     PetscCall(MatSetOption(ab->Be, MAT_ROW_ORIENTED, PETSC_FALSE));
25855d8c7819SPierre Jolivet 
25865d8c7819SPierre Jolivet     /* compute symbolic Ce = Ae*Be */
25875d8c7819SPierre Jolivet     PetscCall(MatCreate(PetscObjectComm((PetscObject)C), &ab->Ce));
25885d8c7819SPierre Jolivet     PetscCall(MatMatMultSymbolic_Elemental(ab->Ae, ab->Be, fill, ab->Ce));
25895d8c7819SPierre Jolivet #else
25905d8c7819SPierre Jolivet     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined");
25915d8c7819SPierre Jolivet #endif
25925d8c7819SPierre Jolivet     break;
25935d8c7819SPierre Jolivet   default: /* alg: "petsc" */
25945d8c7819SPierre Jolivet     ab->Ae = NULL;
25955d8c7819SPierre Jolivet     PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, A->cmap->N, B->cmap->N, NULL, &ab->Be));
25965d8c7819SPierre Jolivet     ab->Ce = NULL;
25975d8c7819SPierre Jolivet     break;
25985d8c7819SPierre Jolivet   }
25996718818eSStefano Zampini 
26006718818eSStefano Zampini   C->product->data       = ab;
26016718818eSStefano Zampini   C->product->destroy    = MatDestroy_MatMatMult_MPIDense_MPIDense;
26024222ddf1SHong Zhang   C->ops->matmultnumeric = MatMatMultNumeric_MPIDense_MPIDense;
26033ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2604320f2790SHong Zhang }
26052ef1f0ffSBarry Smith 
2606d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_AB(Mat C)
2607d71ae5a4SJacob Faibussowitsch {
26085d8c7819SPierre Jolivet   Mat_Product *product     = C->product;
26095d8c7819SPierre Jolivet   const char  *algTypes[2] = {"petsc", "elemental"};
26105d8c7819SPierre Jolivet   PetscInt     alg, nalg = PetscDefined(HAVE_ELEMENTAL) ? 2 : 1;
26115d8c7819SPierre Jolivet   PetscBool    flg = PETSC_FALSE;
26125d8c7819SPierre Jolivet 
26134d86920dSPierre Jolivet   PetscFunctionBegin;
26145d8c7819SPierre Jolivet   /* Set default algorithm */
2615f0b74427SPierre Jolivet   alg = 0; /* default is PETSc */
26165d8c7819SPierre Jolivet   PetscCall(PetscStrcmp(product->alg, "default", &flg));
2617835f2295SStefano Zampini   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
26185d8c7819SPierre Jolivet 
26195d8c7819SPierre Jolivet   /* Get runtime option */
26205d8c7819SPierre Jolivet   PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_AB", "Mat");
26215d8c7819SPierre Jolivet   PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_AB", algTypes, nalg, algTypes[alg], &alg, &flg));
26225d8c7819SPierre Jolivet   PetscOptionsEnd();
2623835f2295SStefano Zampini   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
26245d8c7819SPierre Jolivet 
26254222ddf1SHong Zhang   C->ops->matmultsymbolic = MatMatMultSymbolic_MPIDense_MPIDense;
26264222ddf1SHong Zhang   C->ops->productsymbolic = MatProductSymbolic_AB;
26273ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2628320f2790SHong Zhang }
262986aefd0dSHong Zhang 
2630d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_AtB(Mat C)
2631d71ae5a4SJacob Faibussowitsch {
26324222ddf1SHong Zhang   Mat_Product *product = C->product;
26334222ddf1SHong Zhang   Mat          A = product->A, B = product->B;
26344222ddf1SHong Zhang 
26354222ddf1SHong Zhang   PetscFunctionBegin;
263607e90273SPierre 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 ")",
263707e90273SPierre Jolivet              A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
26386718818eSStefano Zampini   C->ops->transposematmultsymbolic = MatTransposeMatMultSymbolic_MPIDense_MPIDense;
26396718818eSStefano Zampini   C->ops->productsymbolic          = MatProductSymbolic_AtB;
26403ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
26414222ddf1SHong Zhang }
26424222ddf1SHong Zhang 
2643d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_ABt(Mat C)
2644d71ae5a4SJacob Faibussowitsch {
26454222ddf1SHong Zhang   Mat_Product *product     = C->product;
26464222ddf1SHong Zhang   const char  *algTypes[2] = {"allgatherv", "cyclic"};
26474222ddf1SHong Zhang   PetscInt     alg, nalg = 2;
26484222ddf1SHong Zhang   PetscBool    flg = PETSC_FALSE;
26494222ddf1SHong Zhang 
26504222ddf1SHong Zhang   PetscFunctionBegin;
26514222ddf1SHong Zhang   /* Set default algorithm */
26524222ddf1SHong Zhang   alg = 0; /* default is allgatherv */
26539566063dSJacob Faibussowitsch   PetscCall(PetscStrcmp(product->alg, "default", &flg));
2654835f2295SStefano Zampini   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
26554222ddf1SHong Zhang 
26564222ddf1SHong Zhang   /* Get runtime option */
26574222ddf1SHong Zhang   if (product->api_user) {
2658d0609cedSBarry Smith     PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatMatTransposeMult", "Mat");
26599566063dSJacob Faibussowitsch     PetscCall(PetscOptionsEList("-matmattransmult_mpidense_mpidense_via", "Algorithmic approach", "MatMatTransposeMult", algTypes, nalg, algTypes[alg], &alg, &flg));
2660d0609cedSBarry Smith     PetscOptionsEnd();
26614222ddf1SHong Zhang   } else {
2662d0609cedSBarry Smith     PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_ABt", "Mat");
26639566063dSJacob Faibussowitsch     PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_ABt", algTypes, nalg, algTypes[alg], &alg, &flg));
2664d0609cedSBarry Smith     PetscOptionsEnd();
26654222ddf1SHong Zhang   }
2666835f2295SStefano Zampini   if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
26674222ddf1SHong Zhang 
26684222ddf1SHong Zhang   C->ops->mattransposemultsymbolic = MatMatTransposeMultSymbolic_MPIDense_MPIDense;
26694222ddf1SHong Zhang   C->ops->productsymbolic          = MatProductSymbolic_ABt;
26703ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
26714222ddf1SHong Zhang }
26724222ddf1SHong Zhang 
26735d8c7819SPierre Jolivet static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat C)
2674d71ae5a4SJacob Faibussowitsch {
26754222ddf1SHong Zhang   Mat_Product *product = C->product;
26764222ddf1SHong Zhang 
26774222ddf1SHong Zhang   PetscFunctionBegin;
26784222ddf1SHong Zhang   switch (product->type) {
2679d71ae5a4SJacob Faibussowitsch   case MATPRODUCT_AB:
2680d71ae5a4SJacob Faibussowitsch     PetscCall(MatProductSetFromOptions_MPIDense_AB(C));
2681d71ae5a4SJacob Faibussowitsch     break;
2682d71ae5a4SJacob Faibussowitsch   case MATPRODUCT_AtB:
2683d71ae5a4SJacob Faibussowitsch     PetscCall(MatProductSetFromOptions_MPIDense_AtB(C));
2684d71ae5a4SJacob Faibussowitsch     break;
2685d71ae5a4SJacob Faibussowitsch   case MATPRODUCT_ABt:
2686d71ae5a4SJacob Faibussowitsch     PetscCall(MatProductSetFromOptions_MPIDense_ABt(C));
2687d71ae5a4SJacob Faibussowitsch     break;
2688d71ae5a4SJacob Faibussowitsch   default:
2689d71ae5a4SJacob Faibussowitsch     break;
26904222ddf1SHong Zhang   }
26913ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
26924222ddf1SHong Zhang }
2693