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 <petsc/private/deviceimpl.h>
12*d484b384SBoris Martin #include <petsc/private/sfimpl.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 @*/
MatDenseGetLocalMatrix(Mat A,Mat * B)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
MatCopy_MPIDense(Mat A,Mat B,MatStructure s)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
MatShift_MPIDense(Mat A,PetscScalar alpha)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
MatGetRow_MPIDense(Mat A,PetscInt row,PetscInt * nz,PetscInt ** idx,PetscScalar ** v)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
MatRestoreRow_MPIDense(Mat A,PetscInt row,PetscInt * nz,PetscInt ** idx,PetscScalar ** v)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
MatGetDiagonalBlock_MPIDense(Mat A,Mat * a)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
MatSetValues_MPIDense(Mat mat,PetscInt m,const PetscInt idxm[],PetscInt n,const PetscInt idxn[],const PetscScalar v[],InsertMode addv)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
MatGetValues_MPIDense(Mat mat,PetscInt m,const PetscInt idxm[],PetscInt n,const PetscInt idxn[],PetscScalar v[])16666976f2fSJacob Faibussowitsch static PetscErrorCode MatGetValues_MPIDense(Mat mat, PetscInt m, const PetscInt idxm[], PetscInt n, const PetscInt idxn[], PetscScalar v[])
167d71ae5a4SJacob Faibussowitsch {
168b49de8d1SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
169d0f46423SBarry Smith PetscInt i, j, rstart = mat->rmap->rstart, rend = mat->rmap->rend, row;
170b49de8d1SLois Curfman McInnes
1713a40ed3dSBarry Smith PetscFunctionBegin;
172b49de8d1SLois Curfman McInnes for (i = 0; i < m; i++) {
17354c59aa7SJacob Faibussowitsch if (idxm[i] < 0) continue; /* negative row */
17454c59aa7SJacob Faibussowitsch PetscCheck(idxm[i] < mat->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large");
175966bd95aSPierre Jolivet PetscCheck(idxm[i] >= rstart && idxm[i] < rend, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only local values currently supported");
176b49de8d1SLois Curfman McInnes row = idxm[i] - rstart;
177b49de8d1SLois Curfman McInnes for (j = 0; j < n; j++) {
17854c59aa7SJacob Faibussowitsch if (idxn[j] < 0) continue; /* negative column */
17954c59aa7SJacob Faibussowitsch PetscCheck(idxn[j] < mat->cmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Column too large");
1809566063dSJacob Faibussowitsch PetscCall(MatGetValues(mdn->A, 1, &row, 1, &idxn[j], v + i * n + j));
181b49de8d1SLois Curfman McInnes }
1828965ea79SLois Curfman McInnes }
1833ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1848965ea79SLois Curfman McInnes }
1858965ea79SLois Curfman McInnes
MatDenseGetLDA_MPIDense(Mat A,PetscInt * lda)186d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetLDA_MPIDense(Mat A, PetscInt *lda)
187d71ae5a4SJacob Faibussowitsch {
18849a6ff4bSBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data;
18949a6ff4bSBarry Smith
19049a6ff4bSBarry Smith PetscFunctionBegin;
1919566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, lda));
1923ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
19349a6ff4bSBarry Smith }
19449a6ff4bSBarry Smith
MatDenseSetLDA_MPIDense(Mat A,PetscInt lda)195d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseSetLDA_MPIDense(Mat A, PetscInt lda)
196d71ae5a4SJacob Faibussowitsch {
197ad16ce7aSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data;
19847d993e7Ssuyashtn MatType mtype = MATSEQDENSE;
199ad16ce7aSStefano Zampini
200ad16ce7aSStefano Zampini PetscFunctionBegin;
20117359960SJose E. Roman if (!a->A) {
20228b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2039566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(A->rmap));
2049566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(A->cmap));
2059566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
2069566063dSJacob Faibussowitsch PetscCall(MatSetSizes(a->A, A->rmap->n, A->cmap->N, A->rmap->n, A->cmap->N));
207a51cd166SPierre Jolivet #if PetscDefined(HAVE_CUDA)
20847d993e7Ssuyashtn PetscBool iscuda;
2099566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSECUDA, &iscuda));
21047d993e7Ssuyashtn if (iscuda) mtype = MATSEQDENSECUDA;
211a51cd166SPierre Jolivet #elif PetscDefined(HAVE_HIP)
21247d993e7Ssuyashtn PetscBool iship;
21347d993e7Ssuyashtn PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSEHIP, &iship));
21447d993e7Ssuyashtn if (iship) mtype = MATSEQDENSEHIP;
21547d993e7Ssuyashtn #endif
21647d993e7Ssuyashtn PetscCall(MatSetType(a->A, mtype));
21717359960SJose E. Roman }
2189566063dSJacob Faibussowitsch PetscCall(MatDenseSetLDA(a->A, lda));
2193ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
220ad16ce7aSStefano Zampini }
221ad16ce7aSStefano Zampini
MatDenseGetArray_MPIDense(Mat A,PetscScalar ** array)222d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetArray_MPIDense(Mat A, PetscScalar **array)
223d71ae5a4SJacob Faibussowitsch {
224ff14e315SSatish Balay Mat_MPIDense *a = (Mat_MPIDense *)A->data;
225ff14e315SSatish Balay
2263a40ed3dSBarry Smith PetscFunctionBegin;
22728b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2289566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(a->A, array));
2293ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
230ff14e315SSatish Balay }
231ff14e315SSatish Balay
MatDenseGetArrayRead_MPIDense(Mat A,PetscScalar ** array)232f1e99121SPierre Jolivet static PetscErrorCode MatDenseGetArrayRead_MPIDense(Mat A, PetscScalar **array)
233d71ae5a4SJacob Faibussowitsch {
2348572280aSBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data;
2358572280aSBarry Smith
2368572280aSBarry Smith PetscFunctionBegin;
23728b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
238f1e99121SPierre Jolivet PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)array));
2393ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2408572280aSBarry Smith }
2418572280aSBarry Smith
MatDenseGetArrayWrite_MPIDense(Mat A,PetscScalar ** array)242d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetArrayWrite_MPIDense(Mat A, PetscScalar **array)
243d71ae5a4SJacob Faibussowitsch {
2446947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data;
2456947451fSStefano Zampini
2466947451fSStefano Zampini PetscFunctionBegin;
24728b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2489566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayWrite(a->A, array));
2493ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2506947451fSStefano Zampini }
2516947451fSStefano Zampini
MatDensePlaceArray_MPIDense(Mat A,const PetscScalar * array)252d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDensePlaceArray_MPIDense(Mat A, const PetscScalar *array)
253d71ae5a4SJacob Faibussowitsch {
254d3042a70SBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data;
255d3042a70SBarry Smith
256d3042a70SBarry Smith PetscFunctionBegin;
25728b400f6SJacob Faibussowitsch PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
25828b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2599566063dSJacob Faibussowitsch PetscCall(MatDensePlaceArray(a->A, array));
2603ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
261d3042a70SBarry Smith }
262d3042a70SBarry Smith
MatDenseResetArray_MPIDense(Mat A)263d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseResetArray_MPIDense(Mat A)
264d71ae5a4SJacob Faibussowitsch {
265d3042a70SBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data;
266d3042a70SBarry Smith
267d3042a70SBarry Smith PetscFunctionBegin;
26828b400f6SJacob Faibussowitsch PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
26928b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2709566063dSJacob Faibussowitsch PetscCall(MatDenseResetArray(a->A));
2713ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
272d3042a70SBarry Smith }
273d3042a70SBarry Smith
MatDenseReplaceArray_MPIDense(Mat A,const PetscScalar * array)274d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseReplaceArray_MPIDense(Mat A, const PetscScalar *array)
275d71ae5a4SJacob Faibussowitsch {
276d5ea218eSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data;
277d5ea218eSStefano Zampini
278d5ea218eSStefano Zampini PetscFunctionBegin;
27928b400f6SJacob Faibussowitsch PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
28028b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
2819566063dSJacob Faibussowitsch PetscCall(MatDenseReplaceArray(a->A, array));
2823ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
283d5ea218eSStefano Zampini }
284d5ea218eSStefano Zampini
MatCreateSubMatrix_MPIDense(Mat A,IS isrow,IS iscol,MatReuse scall,Mat * B)285d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatCreateSubMatrix_MPIDense(Mat A, IS isrow, IS iscol, MatReuse scall, Mat *B)
286d71ae5a4SJacob Faibussowitsch {
287ca3fa75bSLois Curfman McInnes Mat_MPIDense *mat = (Mat_MPIDense *)A->data, *newmatd;
288637a0070SStefano Zampini PetscInt lda, i, j, rstart, rend, nrows, ncols, Ncols, nlrows, nlcols;
2895d0c19d7SBarry Smith const PetscInt *irow, *icol;
290637a0070SStefano Zampini const PetscScalar *v;
291637a0070SStefano Zampini PetscScalar *bv;
292ca3fa75bSLois Curfman McInnes Mat newmat;
2934aa3045dSJed Brown IS iscol_local;
29442a884f0SBarry Smith MPI_Comm comm_is, comm_mat;
295ca3fa75bSLois Curfman McInnes
296ca3fa75bSLois Curfman McInnes PetscFunctionBegin;
2979566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A, &comm_mat));
2989566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)iscol, &comm_is));
29908401ef6SPierre Jolivet PetscCheck(comm_mat == comm_is, PETSC_COMM_SELF, PETSC_ERR_ARG_NOTSAMECOMM, "IS communicator must match matrix communicator");
30042a884f0SBarry Smith
3019566063dSJacob Faibussowitsch PetscCall(ISAllGather(iscol, &iscol_local));
3029566063dSJacob Faibussowitsch PetscCall(ISGetIndices(isrow, &irow));
3039566063dSJacob Faibussowitsch PetscCall(ISGetIndices(iscol_local, &icol));
3049566063dSJacob Faibussowitsch PetscCall(ISGetLocalSize(isrow, &nrows));
3059566063dSJacob Faibussowitsch PetscCall(ISGetLocalSize(iscol, &ncols));
3069566063dSJacob Faibussowitsch PetscCall(ISGetSize(iscol, &Ncols)); /* global number of columns, size of iscol_local */
307ca3fa75bSLois Curfman McInnes
308ca3fa75bSLois Curfman McInnes /* No parallel redistribution currently supported! Should really check each index set
309da81f932SPierre Jolivet to confirm that it is OK. ... Currently supports only submatrix same partitioning as
3107eba5e9cSLois Curfman McInnes original matrix! */
311ca3fa75bSLois Curfman McInnes
3129566063dSJacob Faibussowitsch PetscCall(MatGetLocalSize(A, &nlrows, &nlcols));
3139566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRange(A, &rstart, &rend));
314ca3fa75bSLois Curfman McInnes
315ca3fa75bSLois Curfman McInnes /* Check submatrix call */
316ca3fa75bSLois Curfman McInnes if (scall == MAT_REUSE_MATRIX) {
317e32f2f54SBarry Smith /* SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Reused submatrix wrong size"); */
3187eba5e9cSLois Curfman McInnes /* Really need to test rows and column sizes! */
319ca3fa75bSLois Curfman McInnes newmat = *B;
320ca3fa75bSLois Curfman McInnes } else {
321ca3fa75bSLois Curfman McInnes /* Create and fill new matrix */
3229566063dSJacob Faibussowitsch PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &newmat));
3239566063dSJacob Faibussowitsch PetscCall(MatSetSizes(newmat, nrows, ncols, PETSC_DECIDE, Ncols));
3249566063dSJacob Faibussowitsch PetscCall(MatSetType(newmat, ((PetscObject)A)->type_name));
3259566063dSJacob Faibussowitsch PetscCall(MatMPIDenseSetPreallocation(newmat, NULL));
326ca3fa75bSLois Curfman McInnes }
327ca3fa75bSLois Curfman McInnes
328ca3fa75bSLois Curfman McInnes /* Now extract the data pointers and do the copy, column at a time */
329ca3fa75bSLois Curfman McInnes newmatd = (Mat_MPIDense *)newmat->data;
3309566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(newmatd->A, &bv));
3319566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(mat->A, &v));
3329566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(mat->A, &lda));
3334aa3045dSJed Brown for (i = 0; i < Ncols; i++) {
334637a0070SStefano Zampini const PetscScalar *av = v + lda * icol[i];
335ad540459SPierre Jolivet for (j = 0; j < nrows; j++) *bv++ = av[irow[j] - rstart];
336ca3fa75bSLois Curfman McInnes }
3379566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(mat->A, &v));
3389566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(newmatd->A, &bv));
339ca3fa75bSLois Curfman McInnes
340ca3fa75bSLois Curfman McInnes /* Assemble the matrices so that the correct flags are set */
3419566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(newmat, MAT_FINAL_ASSEMBLY));
3429566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(newmat, MAT_FINAL_ASSEMBLY));
343ca3fa75bSLois Curfman McInnes
344ca3fa75bSLois Curfman McInnes /* Free work space */
3459566063dSJacob Faibussowitsch PetscCall(ISRestoreIndices(isrow, &irow));
3469566063dSJacob Faibussowitsch PetscCall(ISRestoreIndices(iscol_local, &icol));
3479566063dSJacob Faibussowitsch PetscCall(ISDestroy(&iscol_local));
348ca3fa75bSLois Curfman McInnes *B = newmat;
3493ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
350ca3fa75bSLois Curfman McInnes }
351ca3fa75bSLois Curfman McInnes
MatDenseRestoreArray_MPIDense(Mat A,PetscScalar ** array)35266976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreArray_MPIDense(Mat A, PetscScalar **array)
353d71ae5a4SJacob Faibussowitsch {
35473a71a0fSBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data;
35573a71a0fSBarry Smith
3563a40ed3dSBarry Smith PetscFunctionBegin;
3579566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(a->A, array));
3583ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
359ff14e315SSatish Balay }
360ff14e315SSatish Balay
MatDenseRestoreArrayRead_MPIDense(Mat A,PetscScalar ** array)361f1e99121SPierre Jolivet static PetscErrorCode MatDenseRestoreArrayRead_MPIDense(Mat A, PetscScalar **array)
362d71ae5a4SJacob Faibussowitsch {
3638572280aSBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data;
3648572280aSBarry Smith
3658572280aSBarry Smith PetscFunctionBegin;
366f1e99121SPierre Jolivet PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)array));
3673ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
3688572280aSBarry Smith }
3698572280aSBarry Smith
MatDenseRestoreArrayWrite_MPIDense(Mat A,PetscScalar ** array)37066976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreArrayWrite_MPIDense(Mat A, PetscScalar **array)
371d71ae5a4SJacob Faibussowitsch {
3726947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data;
3736947451fSStefano Zampini
3746947451fSStefano Zampini PetscFunctionBegin;
3759566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayWrite(a->A, array));
3763ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
3776947451fSStefano Zampini }
3786947451fSStefano Zampini
MatAssemblyBegin_MPIDense(Mat mat,MatAssemblyType mode)37966976f2fSJacob Faibussowitsch static PetscErrorCode MatAssemblyBegin_MPIDense(Mat mat, MatAssemblyType mode)
380d71ae5a4SJacob Faibussowitsch {
38139ddd567SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
38213f74950SBarry Smith PetscInt nstash, reallocs;
3838965ea79SLois Curfman McInnes
3843a40ed3dSBarry Smith PetscFunctionBegin;
3853ba16761SJacob Faibussowitsch if (mdn->donotstash || mat->nooffprocentries) PetscFunctionReturn(PETSC_SUCCESS);
3868965ea79SLois Curfman McInnes
3879566063dSJacob Faibussowitsch PetscCall(MatStashScatterBegin_Private(mat, &mat->stash, mat->rmap->range));
3889566063dSJacob Faibussowitsch PetscCall(MatStashGetInfo_Private(&mat->stash, &nstash, &reallocs));
3899566063dSJacob Faibussowitsch PetscCall(PetscInfo(mdn->A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
3903ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
3918965ea79SLois Curfman McInnes }
3928965ea79SLois Curfman McInnes
MatAssemblyEnd_MPIDense(Mat mat,MatAssemblyType mode)39366976f2fSJacob Faibussowitsch static PetscErrorCode MatAssemblyEnd_MPIDense(Mat mat, MatAssemblyType mode)
394d71ae5a4SJacob Faibussowitsch {
39539ddd567SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
39613f74950SBarry Smith PetscInt i, *row, *col, flg, j, rstart, ncols;
39713f74950SBarry Smith PetscMPIInt n;
39887828ca2SBarry Smith PetscScalar *val;
3998965ea79SLois Curfman McInnes
4003a40ed3dSBarry Smith PetscFunctionBegin;
401910cf402Sprj- if (!mdn->donotstash && !mat->nooffprocentries) {
4028965ea79SLois Curfman McInnes /* wait on receives */
4037ef1d9bdSSatish Balay while (1) {
4049566063dSJacob Faibussowitsch PetscCall(MatStashScatterGetMesg_Private(&mat->stash, &n, &row, &col, &val, &flg));
4057ef1d9bdSSatish Balay if (!flg) break;
4068965ea79SLois Curfman McInnes
4077ef1d9bdSSatish Balay for (i = 0; i < n;) {
4087ef1d9bdSSatish Balay /* Now identify the consecutive vals belonging to the same row */
4092205254eSKarl Rupp for (j = i, rstart = row[j]; j < n; j++) {
4102205254eSKarl Rupp if (row[j] != rstart) break;
4112205254eSKarl Rupp }
4127ef1d9bdSSatish Balay if (j < n) ncols = j - i;
4137ef1d9bdSSatish Balay else ncols = n - i;
4147ef1d9bdSSatish Balay /* Now assemble all these values with a single function call */
4159566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIDense(mat, 1, row + i, ncols, col + i, val + i, mat->insertmode));
4167ef1d9bdSSatish Balay i = j;
4178965ea79SLois Curfman McInnes }
4187ef1d9bdSSatish Balay }
4199566063dSJacob Faibussowitsch PetscCall(MatStashScatterEnd_Private(&mat->stash));
420910cf402Sprj- }
4218965ea79SLois Curfman McInnes
4229566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(mdn->A, mode));
4239566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(mdn->A, mode));
4243ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
4258965ea79SLois Curfman McInnes }
4268965ea79SLois Curfman McInnes
MatZeroEntries_MPIDense(Mat A)42766976f2fSJacob Faibussowitsch static PetscErrorCode MatZeroEntries_MPIDense(Mat A)
428d71ae5a4SJacob Faibussowitsch {
42939ddd567SLois Curfman McInnes Mat_MPIDense *l = (Mat_MPIDense *)A->data;
4303a40ed3dSBarry Smith
4313a40ed3dSBarry Smith PetscFunctionBegin;
4329566063dSJacob Faibussowitsch PetscCall(MatZeroEntries(l->A));
4333ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
4348965ea79SLois Curfman McInnes }
4358965ea79SLois Curfman McInnes
MatZeroRows_MPIDense(Mat A,PetscInt n,const PetscInt rows[],PetscScalar diag,Vec x,Vec b)43666976f2fSJacob Faibussowitsch static PetscErrorCode MatZeroRows_MPIDense(Mat A, PetscInt n, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
437d71ae5a4SJacob Faibussowitsch {
43839ddd567SLois Curfman McInnes Mat_MPIDense *l = (Mat_MPIDense *)A->data;
439637a0070SStefano Zampini PetscInt i, len, *lrows;
440637a0070SStefano Zampini
441637a0070SStefano Zampini PetscFunctionBegin;
442637a0070SStefano Zampini /* get locally owned rows */
4439566063dSJacob Faibussowitsch PetscCall(PetscLayoutMapLocal(A->rmap, n, rows, &len, &lrows, NULL));
444dd8e379bSPierre Jolivet /* fix right-hand side if needed */
445637a0070SStefano Zampini if (x && b) {
44697b48c8fSBarry Smith const PetscScalar *xx;
44797b48c8fSBarry Smith PetscScalar *bb;
4488965ea79SLois Curfman McInnes
4499566063dSJacob Faibussowitsch PetscCall(VecGetArrayRead(x, &xx));
4509566063dSJacob Faibussowitsch PetscCall(VecGetArrayWrite(b, &bb));
451637a0070SStefano Zampini for (i = 0; i < len; ++i) bb[lrows[i]] = diag * xx[lrows[i]];
4529566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayRead(x, &xx));
4539566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayWrite(b, &bb));
45497b48c8fSBarry Smith }
4559566063dSJacob Faibussowitsch PetscCall(MatZeroRows(l->A, len, lrows, 0.0, NULL, NULL));
456e2eb51b1SBarry Smith if (diag != 0.0) {
457637a0070SStefano Zampini Vec d;
458b9679d65SBarry Smith
4599566063dSJacob Faibussowitsch PetscCall(MatCreateVecs(A, NULL, &d));
4609566063dSJacob Faibussowitsch PetscCall(VecSet(d, diag));
4619566063dSJacob Faibussowitsch PetscCall(MatDiagonalSet(A, d, INSERT_VALUES));
4629566063dSJacob Faibussowitsch PetscCall(VecDestroy(&d));
463b9679d65SBarry Smith }
4649566063dSJacob Faibussowitsch PetscCall(PetscFree(lrows));
4653ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
4668965ea79SLois Curfman McInnes }
4678965ea79SLois Curfman McInnes
4680be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMult_SeqDense(Mat, Vec, Vec);
4690be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMultAdd_SeqDense(Mat, Vec, Vec, Vec);
4700be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMultTranspose_SeqDense(Mat, Vec, Vec);
4710be0d8bdSHansol Suh PETSC_INTERN PetscErrorCode MatMultTransposeAdd_SeqDense(Mat, Vec, Vec, Vec);
4720be0d8bdSHansol Suh
MatMultColumnRange_MPIDense(Mat mat,Vec xx,Vec yy,PetscInt c_start,PetscInt c_end)473d016bddeSToby Isaac static PetscErrorCode MatMultColumnRange_MPIDense(Mat mat, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end)
474d016bddeSToby Isaac {
475d016bddeSToby Isaac Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
476d016bddeSToby Isaac const PetscScalar *ax;
477d016bddeSToby Isaac PetscScalar *ay;
478d016bddeSToby Isaac PetscMemType axmtype, aymtype;
479d016bddeSToby Isaac
480d016bddeSToby Isaac PetscFunctionBegin;
481d016bddeSToby Isaac if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
482d016bddeSToby Isaac PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
483d016bddeSToby Isaac PetscCall(VecGetArrayWriteAndMemType(mdn->lvec, &ay, &aymtype));
484d016bddeSToby Isaac PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
485d016bddeSToby Isaac PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
486d016bddeSToby Isaac PetscCall(VecRestoreArrayWriteAndMemType(mdn->lvec, &ay));
487d016bddeSToby Isaac PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
488d016bddeSToby Isaac PetscUseMethod(mdn->A, "MatMultColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (mdn->A, mdn->lvec, yy, c_start, c_end));
489d016bddeSToby Isaac PetscFunctionReturn(PETSC_SUCCESS);
490d016bddeSToby Isaac }
491d016bddeSToby Isaac
MatMult_MPIDense(Mat mat,Vec xx,Vec yy)49266976f2fSJacob Faibussowitsch static PetscErrorCode MatMult_MPIDense(Mat mat, Vec xx, Vec yy)
493d71ae5a4SJacob Faibussowitsch {
49439ddd567SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
495637a0070SStefano Zampini const PetscScalar *ax;
496637a0070SStefano Zampini PetscScalar *ay;
497d0295fc0SJunchao Zhang PetscMemType axmtype, aymtype;
498c456f294SBarry Smith
4993a40ed3dSBarry Smith PetscFunctionBegin;
5006f08ad05SPierre Jolivet if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
5019566063dSJacob Faibussowitsch PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
502d016bddeSToby Isaac PetscCall(VecGetArrayWriteAndMemType(mdn->lvec, &ay, &aymtype));
5039566063dSJacob Faibussowitsch PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
5049566063dSJacob Faibussowitsch PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
505d016bddeSToby Isaac PetscCall(VecRestoreArrayWriteAndMemType(mdn->lvec, &ay));
5069566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
5079566063dSJacob Faibussowitsch PetscCall((*mdn->A->ops->mult)(mdn->A, mdn->lvec, yy));
5083ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
5098965ea79SLois Curfman McInnes }
5108965ea79SLois Curfman McInnes
MatMultAddColumnRange_MPIDense(Mat mat,Vec xx,Vec yy,Vec zz,PetscInt c_start,PetscInt c_end)5110be0d8bdSHansol Suh static PetscErrorCode MatMultAddColumnRange_MPIDense(Mat mat, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end)
5120be0d8bdSHansol Suh {
5130be0d8bdSHansol Suh Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
5140be0d8bdSHansol Suh const PetscScalar *ax;
5150be0d8bdSHansol Suh PetscScalar *ay;
5160be0d8bdSHansol Suh PetscMemType axmtype, aymtype;
5170be0d8bdSHansol Suh
5180be0d8bdSHansol Suh PetscFunctionBegin;
5190be0d8bdSHansol Suh if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
5200be0d8bdSHansol Suh PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
5210be0d8bdSHansol Suh PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype));
5220be0d8bdSHansol Suh PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
5230be0d8bdSHansol Suh PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
5240be0d8bdSHansol Suh PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay));
5250be0d8bdSHansol Suh PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
5260be0d8bdSHansol Suh PetscUseMethod(mdn->A, "MatMultAddColumnRange_C", (Mat, Vec, Vec, Vec, PetscInt, PetscInt), (mdn->A, mdn->lvec, yy, zz, c_start, c_end));
5270be0d8bdSHansol Suh PetscFunctionReturn(PETSC_SUCCESS);
5280be0d8bdSHansol Suh }
5290be0d8bdSHansol Suh
MatMultAdd_MPIDense(Mat mat,Vec xx,Vec yy,Vec zz)53066976f2fSJacob Faibussowitsch static PetscErrorCode MatMultAdd_MPIDense(Mat mat, Vec xx, Vec yy, Vec zz)
531d71ae5a4SJacob Faibussowitsch {
53239ddd567SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
533637a0070SStefano Zampini const PetscScalar *ax;
534637a0070SStefano Zampini PetscScalar *ay;
535d0295fc0SJunchao Zhang PetscMemType axmtype, aymtype;
536c456f294SBarry Smith
5373a40ed3dSBarry Smith PetscFunctionBegin;
5386f08ad05SPierre Jolivet if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
5399566063dSJacob Faibussowitsch PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
5409566063dSJacob Faibussowitsch PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype));
5419566063dSJacob Faibussowitsch PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
5429566063dSJacob Faibussowitsch PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
5439566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay));
5449566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
5459566063dSJacob Faibussowitsch PetscCall((*mdn->A->ops->multadd)(mdn->A, mdn->lvec, yy, zz));
5463ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
5478965ea79SLois Curfman McInnes }
5488965ea79SLois Curfman McInnes
MatMultHermitianTransposeColumnRange_MPIDense(Mat A,Vec xx,Vec yy,PetscInt c_start,PetscInt c_end)5490be0d8bdSHansol Suh static PetscErrorCode MatMultHermitianTransposeColumnRange_MPIDense(Mat A, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end)
5500be0d8bdSHansol Suh {
5510be0d8bdSHansol Suh Mat_MPIDense *a = (Mat_MPIDense *)A->data;
5520be0d8bdSHansol Suh const PetscScalar *ax;
5530be0d8bdSHansol Suh PetscScalar *ay;
5540be0d8bdSHansol Suh PetscMemType axmtype, aymtype;
555d016bddeSToby Isaac PetscInt r_start, r_end;
556d016bddeSToby Isaac PetscInt c_start_local, c_end_local;
5570be0d8bdSHansol Suh
5580be0d8bdSHansol Suh PetscFunctionBegin;
5590be0d8bdSHansol Suh if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
560d016bddeSToby Isaac PetscCall(VecZeroEntries(a->lvec));
561d016bddeSToby Isaac PetscCall(VecGetOwnershipRange(yy, &r_start, &r_end));
562d016bddeSToby Isaac c_start_local = PetscMax(c_start, r_start);
563d016bddeSToby Isaac c_end_local = PetscMin(c_end, r_end);
564d016bddeSToby Isaac PetscCall(VecGetArrayAndMemType(yy, &ay, &aymtype));
565d016bddeSToby Isaac if (c_end_local > c_start_local) {
566d016bddeSToby Isaac if (PetscMemTypeHost(aymtype)) {
567d016bddeSToby Isaac PetscCall(PetscArrayzero(&ay[c_start_local], (size_t)(c_end_local - c_start_local)));
568d016bddeSToby Isaac } else {
569d016bddeSToby Isaac PetscCall(PetscDeviceRegisterMemory(ay, aymtype, sizeof(*ay) * ((size_t)(r_end - r_start))));
570d016bddeSToby Isaac PetscCall(PetscDeviceArrayZero(NULL, &ay[c_start_local], (size_t)(c_end_local - c_start_local)));
571d016bddeSToby Isaac }
572d016bddeSToby Isaac }
5730be0d8bdSHansol Suh PetscUseMethod(a->A, "MatMultHermitianTransposeColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (a->A, xx, a->lvec, c_start, c_end));
5740be0d8bdSHansol Suh PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
5750be0d8bdSHansol Suh PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
5760be0d8bdSHansol Suh PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
5770be0d8bdSHansol Suh PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
5780be0d8bdSHansol Suh PetscCall(VecRestoreArrayAndMemType(yy, &ay));
5790be0d8bdSHansol Suh PetscFunctionReturn(PETSC_SUCCESS);
5800be0d8bdSHansol Suh }
5810be0d8bdSHansol Suh
MatMultTransposeKernel_MPIDense(Mat A,Vec xx,Vec yy,PetscBool herm)582459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTransposeKernel_MPIDense(Mat A, Vec xx, Vec yy, PetscBool herm)
583d71ae5a4SJacob Faibussowitsch {
584096963f5SLois Curfman McInnes Mat_MPIDense *a = (Mat_MPIDense *)A->data;
585637a0070SStefano Zampini const PetscScalar *ax;
586637a0070SStefano Zampini PetscScalar *ay;
587d0295fc0SJunchao Zhang PetscMemType axmtype, aymtype;
5884d86920dSPierre Jolivet
5893a40ed3dSBarry Smith PetscFunctionBegin;
5906f08ad05SPierre Jolivet if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
5919566063dSJacob Faibussowitsch PetscCall(VecSet(yy, 0.0));
592459e8d23SBlanca Mellado Pinto if (herm) PetscCall((*a->A->ops->multhermitiantranspose)(a->A, xx, a->lvec));
593459e8d23SBlanca Mellado Pinto else PetscCall((*a->A->ops->multtranspose)(a->A, xx, a->lvec));
5949566063dSJacob Faibussowitsch PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
5959566063dSJacob Faibussowitsch PetscCall(VecGetArrayAndMemType(yy, &ay, &aymtype));
5969566063dSJacob Faibussowitsch PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
5979566063dSJacob Faibussowitsch PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
5989566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
5999566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayAndMemType(yy, &ay));
6003ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
601096963f5SLois Curfman McInnes }
602096963f5SLois Curfman McInnes
MatMultHermitianTransposeAddColumnRange_MPIDense(Mat A,Vec xx,Vec yy,Vec zz,PetscInt c_start,PetscInt c_end)6030be0d8bdSHansol Suh static PetscErrorCode MatMultHermitianTransposeAddColumnRange_MPIDense(Mat A, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end)
6040be0d8bdSHansol Suh {
6050be0d8bdSHansol Suh Mat_MPIDense *a = (Mat_MPIDense *)A->data;
6060be0d8bdSHansol Suh const PetscScalar *ax;
6070be0d8bdSHansol Suh PetscScalar *ay;
6080be0d8bdSHansol Suh PetscMemType axmtype, aymtype;
6090be0d8bdSHansol Suh
6100be0d8bdSHansol Suh PetscFunctionBegin;
6110be0d8bdSHansol Suh if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
6120be0d8bdSHansol Suh PetscCall(VecCopy(yy, zz));
613d016bddeSToby Isaac PetscCall(VecZeroEntries(a->lvec));
6140be0d8bdSHansol Suh PetscUseMethod(a->A, "MatMultHermitianTransposeColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (a->A, xx, a->lvec, c_start, c_end));
6150be0d8bdSHansol Suh PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
6160be0d8bdSHansol Suh PetscCall(VecGetArrayAndMemType(zz, &ay, &aymtype));
6170be0d8bdSHansol Suh PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
6180be0d8bdSHansol Suh PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
6190be0d8bdSHansol Suh PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
6200be0d8bdSHansol Suh PetscCall(VecRestoreArrayAndMemType(zz, &ay));
6210be0d8bdSHansol Suh PetscFunctionReturn(PETSC_SUCCESS);
6220be0d8bdSHansol Suh }
6230be0d8bdSHansol Suh
MatMultTransposeAddKernel_MPIDense(Mat A,Vec xx,Vec yy,Vec zz,PetscBool herm)624459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTransposeAddKernel_MPIDense(Mat A, Vec xx, Vec yy, Vec zz, PetscBool herm)
625d71ae5a4SJacob Faibussowitsch {
626096963f5SLois Curfman McInnes Mat_MPIDense *a = (Mat_MPIDense *)A->data;
627637a0070SStefano Zampini const PetscScalar *ax;
628637a0070SStefano Zampini PetscScalar *ay;
629d0295fc0SJunchao Zhang PetscMemType axmtype, aymtype;
630096963f5SLois Curfman McInnes
6313a40ed3dSBarry Smith PetscFunctionBegin;
6326f08ad05SPierre Jolivet if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
6339566063dSJacob Faibussowitsch PetscCall(VecCopy(yy, zz));
634459e8d23SBlanca Mellado Pinto if (herm) PetscCall((*a->A->ops->multhermitiantranspose)(a->A, xx, a->lvec));
635459e8d23SBlanca Mellado Pinto else PetscCall((*a->A->ops->multtranspose)(a->A, xx, a->lvec));
6369566063dSJacob Faibussowitsch PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
6379566063dSJacob Faibussowitsch PetscCall(VecGetArrayAndMemType(zz, &ay, &aymtype));
6389566063dSJacob Faibussowitsch PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
6399566063dSJacob Faibussowitsch PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
6409566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
6419566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayAndMemType(zz, &ay));
6423ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
643096963f5SLois Curfman McInnes }
644096963f5SLois Curfman McInnes
MatMultTranspose_MPIDense(Mat A,Vec xx,Vec yy)645459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTranspose_MPIDense(Mat A, Vec xx, Vec yy)
646459e8d23SBlanca Mellado Pinto {
647459e8d23SBlanca Mellado Pinto PetscFunctionBegin;
648459e8d23SBlanca Mellado Pinto PetscCall(MatMultTransposeKernel_MPIDense(A, xx, yy, PETSC_FALSE));
649459e8d23SBlanca Mellado Pinto PetscFunctionReturn(PETSC_SUCCESS);
650459e8d23SBlanca Mellado Pinto }
651459e8d23SBlanca Mellado Pinto
MatMultTransposeAdd_MPIDense(Mat A,Vec xx,Vec yy,Vec zz)652459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultTransposeAdd_MPIDense(Mat A, Vec xx, Vec yy, Vec zz)
653459e8d23SBlanca Mellado Pinto {
654459e8d23SBlanca Mellado Pinto PetscFunctionBegin;
655459e8d23SBlanca Mellado Pinto PetscCall(MatMultTransposeAddKernel_MPIDense(A, xx, yy, zz, PETSC_FALSE));
656459e8d23SBlanca Mellado Pinto PetscFunctionReturn(PETSC_SUCCESS);
657459e8d23SBlanca Mellado Pinto }
658459e8d23SBlanca Mellado Pinto
MatMultHermitianTranspose_MPIDense(Mat A,Vec xx,Vec yy)659459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultHermitianTranspose_MPIDense(Mat A, Vec xx, Vec yy)
660459e8d23SBlanca Mellado Pinto {
661459e8d23SBlanca Mellado Pinto PetscFunctionBegin;
662459e8d23SBlanca Mellado Pinto PetscCall(MatMultTransposeKernel_MPIDense(A, xx, yy, PETSC_TRUE));
663459e8d23SBlanca Mellado Pinto PetscFunctionReturn(PETSC_SUCCESS);
664459e8d23SBlanca Mellado Pinto }
665459e8d23SBlanca Mellado Pinto
MatMultHermitianTransposeAdd_MPIDense(Mat A,Vec xx,Vec yy,Vec zz)666459e8d23SBlanca Mellado Pinto static PetscErrorCode MatMultHermitianTransposeAdd_MPIDense(Mat A, Vec xx, Vec yy, Vec zz)
667459e8d23SBlanca Mellado Pinto {
668459e8d23SBlanca Mellado Pinto PetscFunctionBegin;
669459e8d23SBlanca Mellado Pinto PetscCall(MatMultTransposeAddKernel_MPIDense(A, xx, yy, zz, PETSC_TRUE));
670459e8d23SBlanca Mellado Pinto PetscFunctionReturn(PETSC_SUCCESS);
671459e8d23SBlanca Mellado Pinto }
672459e8d23SBlanca Mellado Pinto
MatGetDiagonal_MPIDense(Mat A,Vec v)673d71ae5a4SJacob Faibussowitsch PetscErrorCode MatGetDiagonal_MPIDense(Mat A, Vec v)
674d71ae5a4SJacob Faibussowitsch {
67539ddd567SLois Curfman McInnes Mat_MPIDense *a = (Mat_MPIDense *)A->data;
67679560b15SJacob Faibussowitsch PetscInt lda, len, i, nl, ng, m = A->rmap->n, radd;
67779560b15SJacob Faibussowitsch PetscScalar *x;
678637a0070SStefano Zampini const PetscScalar *av;
679ed3cc1f0SBarry Smith
6803a40ed3dSBarry Smith PetscFunctionBegin;
6819566063dSJacob Faibussowitsch PetscCall(VecGetArray(v, &x));
68279560b15SJacob Faibussowitsch PetscCall(VecGetSize(v, &ng));
68379560b15SJacob Faibussowitsch PetscCheck(ng == A->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Nonconforming mat and vec");
68479560b15SJacob Faibussowitsch PetscCall(VecGetLocalSize(v, &nl));
685d0f46423SBarry Smith len = PetscMin(a->A->rmap->n, a->A->cmap->n);
686d0f46423SBarry Smith radd = A->rmap->rstart * m;
6879566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(a->A, &av));
6889566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &lda));
689ad540459SPierre Jolivet for (i = 0; i < len; i++) x[i] = av[radd + i * lda + i];
6909566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(a->A, &av));
6910be0d8bdSHansol Suh if (nl - i > 0) PetscCall(PetscArrayzero(x + i, nl - i));
6929566063dSJacob Faibussowitsch PetscCall(VecRestoreArray(v, &x));
6933ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
6948965ea79SLois Curfman McInnes }
6958965ea79SLois Curfman McInnes
MatDestroy_MPIDense(Mat mat)69666976f2fSJacob Faibussowitsch static PetscErrorCode MatDestroy_MPIDense(Mat mat)
697d71ae5a4SJacob Faibussowitsch {
6983501a2bdSLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
699ed3cc1f0SBarry Smith
7003a40ed3dSBarry Smith PetscFunctionBegin;
7013ba16761SJacob Faibussowitsch PetscCall(PetscLogObjectState((PetscObject)mat, "Rows=%" PetscInt_FMT ", Cols=%" PetscInt_FMT, mat->rmap->N, mat->cmap->N));
7029566063dSJacob Faibussowitsch PetscCall(MatStashDestroy_Private(&mat->stash));
70328b400f6SJacob Faibussowitsch PetscCheck(!mdn->vecinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
70428b400f6SJacob Faibussowitsch PetscCheck(!mdn->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
7059566063dSJacob Faibussowitsch PetscCall(MatDestroy(&mdn->A));
7069566063dSJacob Faibussowitsch PetscCall(VecDestroy(&mdn->lvec));
7079566063dSJacob Faibussowitsch PetscCall(PetscSFDestroy(&mdn->Mvctx));
7089566063dSJacob Faibussowitsch PetscCall(VecDestroy(&mdn->cvec));
7099566063dSJacob Faibussowitsch PetscCall(MatDestroy(&mdn->cmat));
71001b82886SBarry Smith
7119566063dSJacob Faibussowitsch PetscCall(PetscFree(mat->data));
7129566063dSJacob Faibussowitsch PetscCall(PetscObjectChangeTypeName((PetscObject)mat, NULL));
7138baccfbdSHong Zhang
7149566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", NULL));
7159566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", NULL));
7169566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", NULL));
7179566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", NULL));
7189566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", NULL));
7199566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", NULL));
7209566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", NULL));
7219566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", NULL));
7229566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", NULL));
7239566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", NULL));
7249566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", NULL));
7259566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", NULL));
7269566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", NULL));
7278baccfbdSHong Zhang #if defined(PETSC_HAVE_ELEMENTAL)
7289566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", NULL));
7298baccfbdSHong Zhang #endif
730d1a032dbSPierre Jolivet #if defined(PETSC_HAVE_SCALAPACK) && (defined(PETSC_USE_REAL_SINGLE) || defined(PETSC_USE_REAL_DOUBLE))
7319566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", NULL));
732d24d4204SJose E. Roman #endif
7339566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", NULL));
7349566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", NULL));
7359566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", NULL));
7366718818eSStefano Zampini #if defined(PETSC_HAVE_CUDA)
7379566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", NULL));
7389566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", NULL));
7392e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", NULL));
7402e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidensecuda_mpidense_C", NULL));
7412e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidensecuda_C", NULL));
7422e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidensecuda_C", NULL));
7432e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensecuda_mpiaij_C", NULL));
7442e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensecuda_mpiaijcusparse_C", NULL));
7452e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArray_C", NULL));
7462e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArrayRead_C", NULL));
7472e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArrayWrite_C", NULL));
7482e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArray_C", NULL));
7492e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArrayRead_C", NULL));
7502e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArrayWrite_C", NULL));
7512e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAPlaceArray_C", NULL));
7522e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAResetArray_C", NULL));
7532e956fe4SStefano Zampini PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAReplaceArray_C", NULL));
7543d9668e3SJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDASetPreallocation_C", NULL));
7556718818eSStefano Zampini #endif
75647d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
75747d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", NULL));
75847d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", NULL));
75947d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", NULL));
76047d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidensehip_mpidense_C", NULL));
76147d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidensehip_C", NULL));
76247d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidensehip_C", NULL));
76347d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensehip_mpiaij_C", NULL));
76447d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensehip_mpiaijhipsparse_C", NULL));
76547d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArray_C", NULL));
76647d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArrayRead_C", NULL));
76747d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArrayWrite_C", NULL));
76847d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArray_C", NULL));
76947d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArrayRead_C", NULL));
77047d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArrayWrite_C", NULL));
77147d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPPlaceArray_C", NULL));
77247d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPResetArray_C", NULL));
77347d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPReplaceArray_C", NULL));
7743d9668e3SJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPSetPreallocation_C", NULL));
77547d993e7Ssuyashtn #endif
7769566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", NULL));
7779566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", NULL));
7789566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", NULL));
7799566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", NULL));
7809566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", NULL));
7819566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", NULL));
7829566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", NULL));
7839566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", NULL));
7849566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", NULL));
7859566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", NULL));
786d016bddeSToby Isaac PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultColumnRange_C", NULL));
7870be0d8bdSHansol Suh PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", NULL));
7880be0d8bdSHansol Suh PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", NULL));
7890be0d8bdSHansol Suh PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", NULL));
790f8103e6bSStefano Zampini
791f8103e6bSStefano Zampini PetscCall(PetscObjectCompose((PetscObject)mat, "DiagonalBlock", NULL));
7923ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
7938965ea79SLois Curfman McInnes }
79439ddd567SLois Curfman McInnes
7959804daf3SBarry Smith #include <petscdraw.h>
MatView_MPIDense_ASCIIorDraworSocket(Mat mat,PetscViewer viewer)796d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatView_MPIDense_ASCIIorDraworSocket(Mat mat, PetscViewer viewer)
797d71ae5a4SJacob Faibussowitsch {
79839ddd567SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
799616b8fbbSStefano Zampini PetscMPIInt rank;
80019fd82e9SBarry Smith PetscViewerType vtype;
8019f196a02SMartin Diehl PetscBool isascii, isdraw;
802b0a32e0cSBarry Smith PetscViewer sviewer;
803f3ef73ceSBarry Smith PetscViewerFormat format;
8048965ea79SLois Curfman McInnes
8053a40ed3dSBarry Smith PetscFunctionBegin;
8069566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mat), &rank));
8079f196a02SMartin Diehl PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
8089566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
8099f196a02SMartin Diehl if (isascii) {
8109566063dSJacob Faibussowitsch PetscCall(PetscViewerGetType(viewer, &vtype));
8119566063dSJacob Faibussowitsch PetscCall(PetscViewerGetFormat(viewer, &format));
812456192e2SBarry Smith if (format == PETSC_VIEWER_ASCII_INFO_DETAIL) {
8134e220ebcSLois Curfman McInnes MatInfo info;
8149566063dSJacob Faibussowitsch PetscCall(MatGetInfo(mat, MAT_LOCAL, &info));
8159566063dSJacob Faibussowitsch PetscCall(PetscViewerASCIIPushSynchronized(viewer));
8169371c9d4SSatish Balay PetscCall(PetscViewerASCIISynchronizedPrintf(viewer, " [%d] local rows %" PetscInt_FMT " nz %" PetscInt_FMT " nz alloced %" PetscInt_FMT " mem %" PetscInt_FMT " \n", rank, mat->rmap->n, (PetscInt)info.nz_used, (PetscInt)info.nz_allocated,
8179371c9d4SSatish Balay (PetscInt)info.memory));
8189566063dSJacob Faibussowitsch PetscCall(PetscViewerFlush(viewer));
8199566063dSJacob Faibussowitsch PetscCall(PetscViewerASCIIPopSynchronized(viewer));
8206f08ad05SPierre Jolivet if (mdn->Mvctx) PetscCall(PetscSFView(mdn->Mvctx, viewer));
8213ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
822fb9695e5SSatish Balay } else if (format == PETSC_VIEWER_ASCII_INFO) {
8233ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
8248965ea79SLois Curfman McInnes }
825f1af5d2fSBarry Smith } else if (isdraw) {
826b0a32e0cSBarry Smith PetscDraw draw;
827ace3abfcSBarry Smith PetscBool isnull;
828f1af5d2fSBarry Smith
8299566063dSJacob Faibussowitsch PetscCall(PetscViewerDrawGetDraw(viewer, 0, &draw));
8309566063dSJacob Faibussowitsch PetscCall(PetscDrawIsNull(draw, &isnull));
8313ba16761SJacob Faibussowitsch if (isnull) PetscFunctionReturn(PETSC_SUCCESS);
832f1af5d2fSBarry Smith }
83377ed5343SBarry Smith
8347da1fb6eSBarry Smith {
8358965ea79SLois Curfman McInnes /* assemble the entire matrix onto first processor. */
8368965ea79SLois Curfman McInnes Mat A;
837d0f46423SBarry Smith PetscInt M = mat->rmap->N, N = mat->cmap->N, m, row, i, nz;
838ba8c8a56SBarry Smith PetscInt *cols;
839ba8c8a56SBarry Smith PetscScalar *vals;
8408965ea79SLois Curfman McInnes
8419566063dSJacob Faibussowitsch PetscCall(MatCreate(PetscObjectComm((PetscObject)mat), &A));
842dd400576SPatrick Sanan if (rank == 0) {
8439566063dSJacob Faibussowitsch PetscCall(MatSetSizes(A, M, N, M, N));
8443a40ed3dSBarry Smith } else {
8459566063dSJacob Faibussowitsch PetscCall(MatSetSizes(A, 0, 0, M, N));
8468965ea79SLois Curfman McInnes }
8477adad957SLisandro Dalcin /* Since this is a temporary matrix, MATMPIDENSE instead of ((PetscObject)A)->type_name here is probably acceptable. */
8489566063dSJacob Faibussowitsch PetscCall(MatSetType(A, MATMPIDENSE));
8499566063dSJacob Faibussowitsch PetscCall(MatMPIDenseSetPreallocation(A, NULL));
8508965ea79SLois Curfman McInnes
85139ddd567SLois Curfman McInnes /* Copy the matrix ... This isn't the most efficient means,
85239ddd567SLois Curfman McInnes but it's quick for now */
85351022da4SBarry Smith A->insertmode = INSERT_VALUES;
8542205254eSKarl Rupp
8552205254eSKarl Rupp row = mat->rmap->rstart;
8562205254eSKarl Rupp m = mdn->A->rmap->n;
8578965ea79SLois Curfman McInnes for (i = 0; i < m; i++) {
8589566063dSJacob Faibussowitsch PetscCall(MatGetRow_MPIDense(mat, row, &nz, &cols, &vals));
8599566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIDense(A, 1, &row, nz, cols, vals, INSERT_VALUES));
8609566063dSJacob Faibussowitsch PetscCall(MatRestoreRow_MPIDense(mat, row, &nz, &cols, &vals));
86139ddd567SLois Curfman McInnes row++;
8628965ea79SLois Curfman McInnes }
8638965ea79SLois Curfman McInnes
8649566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
8659566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
8669566063dSJacob Faibussowitsch PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
867dd400576SPatrick Sanan if (rank == 0) {
868f4f49eeaSPierre Jolivet PetscCall(PetscObjectSetName((PetscObject)((Mat_MPIDense *)A->data)->A, ((PetscObject)mat)->name));
869f4f49eeaSPierre Jolivet PetscCall(MatView_SeqDense(((Mat_MPIDense *)A->data)->A, sviewer));
8708965ea79SLois Curfman McInnes }
8719566063dSJacob Faibussowitsch PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
8729566063dSJacob Faibussowitsch PetscCall(MatDestroy(&A));
8738965ea79SLois Curfman McInnes }
8743ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
8758965ea79SLois Curfman McInnes }
8768965ea79SLois Curfman McInnes
MatView_MPIDense(Mat mat,PetscViewer viewer)87766976f2fSJacob Faibussowitsch static PetscErrorCode MatView_MPIDense(Mat mat, PetscViewer viewer)
878d71ae5a4SJacob Faibussowitsch {
8799f196a02SMartin Diehl PetscBool isascii, isbinary, isdraw, issocket;
8808965ea79SLois Curfman McInnes
881433994e6SBarry Smith PetscFunctionBegin;
8829f196a02SMartin Diehl PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
8839566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
8849566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERSOCKET, &issocket));
8859566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
8860f5bd95cSBarry Smith
887ac530a7eSPierre Jolivet if (isascii || issocket || isdraw) PetscCall(MatView_MPIDense_ASCIIorDraworSocket(mat, viewer));
888ac530a7eSPierre Jolivet else if (isbinary) PetscCall(MatView_Dense_Binary(mat, viewer));
8893ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
8908965ea79SLois Curfman McInnes }
8918965ea79SLois Curfman McInnes
MatGetInfo_MPIDense(Mat A,MatInfoType flag,MatInfo * info)89266976f2fSJacob Faibussowitsch static PetscErrorCode MatGetInfo_MPIDense(Mat A, MatInfoType flag, MatInfo *info)
893d71ae5a4SJacob Faibussowitsch {
8943501a2bdSLois Curfman McInnes Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
8953501a2bdSLois Curfman McInnes Mat mdn = mat->A;
8963966268fSBarry Smith PetscLogDouble isend[5], irecv[5];
8978965ea79SLois Curfman McInnes
8983a40ed3dSBarry Smith PetscFunctionBegin;
8994e220ebcSLois Curfman McInnes info->block_size = 1.0;
9002205254eSKarl Rupp
9019566063dSJacob Faibussowitsch PetscCall(MatGetInfo(mdn, MAT_LOCAL, info));
9022205254eSKarl Rupp
9039371c9d4SSatish Balay isend[0] = info->nz_used;
9049371c9d4SSatish Balay isend[1] = info->nz_allocated;
9059371c9d4SSatish Balay isend[2] = info->nz_unneeded;
9069371c9d4SSatish Balay isend[3] = info->memory;
9079371c9d4SSatish Balay isend[4] = info->mallocs;
9088965ea79SLois Curfman McInnes if (flag == MAT_LOCAL) {
9094e220ebcSLois Curfman McInnes info->nz_used = isend[0];
9104e220ebcSLois Curfman McInnes info->nz_allocated = isend[1];
9114e220ebcSLois Curfman McInnes info->nz_unneeded = isend[2];
9124e220ebcSLois Curfman McInnes info->memory = isend[3];
9134e220ebcSLois Curfman McInnes info->mallocs = isend[4];
9148965ea79SLois Curfman McInnes } else if (flag == MAT_GLOBAL_MAX) {
915462c564dSBarry Smith PetscCallMPI(MPIU_Allreduce(isend, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_MAX, PetscObjectComm((PetscObject)A)));
9162205254eSKarl Rupp
9174e220ebcSLois Curfman McInnes info->nz_used = irecv[0];
9184e220ebcSLois Curfman McInnes info->nz_allocated = irecv[1];
9194e220ebcSLois Curfman McInnes info->nz_unneeded = irecv[2];
9204e220ebcSLois Curfman McInnes info->memory = irecv[3];
9214e220ebcSLois Curfman McInnes info->mallocs = irecv[4];
9228965ea79SLois Curfman McInnes } else if (flag == MAT_GLOBAL_SUM) {
923462c564dSBarry Smith PetscCallMPI(MPIU_Allreduce(isend, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_SUM, PetscObjectComm((PetscObject)A)));
9242205254eSKarl Rupp
9254e220ebcSLois Curfman McInnes info->nz_used = irecv[0];
9264e220ebcSLois Curfman McInnes info->nz_allocated = irecv[1];
9274e220ebcSLois Curfman McInnes info->nz_unneeded = irecv[2];
9284e220ebcSLois Curfman McInnes info->memory = irecv[3];
9294e220ebcSLois Curfman McInnes info->mallocs = irecv[4];
9308965ea79SLois Curfman McInnes }
9314e220ebcSLois Curfman McInnes info->fill_ratio_given = 0; /* no parallel LU/ILU/Cholesky */
9324e220ebcSLois Curfman McInnes info->fill_ratio_needed = 0;
9334e220ebcSLois Curfman McInnes info->factor_mallocs = 0;
9343ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
9358965ea79SLois Curfman McInnes }
9368965ea79SLois Curfman McInnes
MatSetOption_MPIDense(Mat A,MatOption op,PetscBool flg)93766976f2fSJacob Faibussowitsch static PetscErrorCode MatSetOption_MPIDense(Mat A, MatOption op, PetscBool flg)
938d71ae5a4SJacob Faibussowitsch {
93939ddd567SLois Curfman McInnes Mat_MPIDense *a = (Mat_MPIDense *)A->data;
9408965ea79SLois Curfman McInnes
9413a40ed3dSBarry Smith PetscFunctionBegin;
94212c028f9SKris Buschelman switch (op) {
943512a5fc5SBarry Smith case MAT_NEW_NONZERO_LOCATIONS:
94412c028f9SKris Buschelman case MAT_NEW_NONZERO_LOCATION_ERR:
94512c028f9SKris Buschelman case MAT_NEW_NONZERO_ALLOCATION_ERR:
94643674050SBarry Smith MatCheckPreallocated(A, 1);
9479566063dSJacob Faibussowitsch PetscCall(MatSetOption(a->A, op, flg));
94812c028f9SKris Buschelman break;
94912c028f9SKris Buschelman case MAT_ROW_ORIENTED:
95043674050SBarry Smith MatCheckPreallocated(A, 1);
9514e0d8c25SBarry Smith a->roworiented = flg;
9529566063dSJacob Faibussowitsch PetscCall(MatSetOption(a->A, op, flg));
95312c028f9SKris Buschelman break;
954d71ae5a4SJacob Faibussowitsch case MAT_IGNORE_OFF_PROC_ENTRIES:
955d71ae5a4SJacob Faibussowitsch a->donotstash = flg;
956d71ae5a4SJacob Faibussowitsch break;
95777e54ba9SKris Buschelman case MAT_SYMMETRIC:
95877e54ba9SKris Buschelman case MAT_STRUCTURALLY_SYMMETRIC:
9599a4540c5SBarry Smith case MAT_HERMITIAN:
9609a4540c5SBarry Smith case MAT_SYMMETRY_ETERNAL:
961b94d7dedSBarry Smith case MAT_STRUCTURAL_SYMMETRY_ETERNAL:
962b94d7dedSBarry Smith case MAT_SPD:
963b94d7dedSBarry Smith case MAT_SPD_ETERNAL:
964b94d7dedSBarry Smith /* if the diagonal matrix is square it inherits some of the properties above */
9659725ddffSPierre Jolivet if (a->A && A->rmap->n == A->cmap->n) PetscCall(MatSetOption(a->A, op, flg));
96677e54ba9SKris Buschelman break;
967d71ae5a4SJacob Faibussowitsch default:
968888c827cSStefano Zampini break;
9693a40ed3dSBarry Smith }
9703ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
9718965ea79SLois Curfman McInnes }
9728965ea79SLois Curfman McInnes
MatDiagonalScale_MPIDense(Mat A,Vec ll,Vec rr)97366976f2fSJacob Faibussowitsch static PetscErrorCode MatDiagonalScale_MPIDense(Mat A, Vec ll, Vec rr)
974d71ae5a4SJacob Faibussowitsch {
9755b2fa520SLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)A->data;
976637a0070SStefano Zampini const PetscScalar *l;
977637a0070SStefano Zampini PetscScalar x, *v, *vv, *r;
978637a0070SStefano Zampini PetscInt i, j, s2a, s3a, s2, s3, m = mdn->A->rmap->n, n = mdn->A->cmap->n, lda;
9795b2fa520SLois Curfman McInnes
9805b2fa520SLois Curfman McInnes PetscFunctionBegin;
9819566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(mdn->A, &vv));
9829566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(mdn->A, &lda));
9839566063dSJacob Faibussowitsch PetscCall(MatGetLocalSize(A, &s2, &s3));
9845b2fa520SLois Curfman McInnes if (ll) {
9859566063dSJacob Faibussowitsch PetscCall(VecGetLocalSize(ll, &s2a));
98608401ef6SPierre Jolivet PetscCheck(s2a == s2, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Left scaling vector non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT, s2a, s2);
9879566063dSJacob Faibussowitsch PetscCall(VecGetArrayRead(ll, &l));
9885b2fa520SLois Curfman McInnes for (i = 0; i < m; i++) {
9895b2fa520SLois Curfman McInnes x = l[i];
990637a0070SStefano Zampini v = vv + i;
9919371c9d4SSatish Balay for (j = 0; j < n; j++) {
9929371c9d4SSatish Balay (*v) *= x;
9939371c9d4SSatish Balay v += lda;
9949371c9d4SSatish Balay }
9955b2fa520SLois Curfman McInnes }
9969566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayRead(ll, &l));
9979566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(1.0 * n * m));
9985b2fa520SLois Curfman McInnes }
9995b2fa520SLois Curfman McInnes if (rr) {
1000637a0070SStefano Zampini const PetscScalar *ar;
1001637a0070SStefano Zampini
10029566063dSJacob Faibussowitsch PetscCall(VecGetLocalSize(rr, &s3a));
100308401ef6SPierre Jolivet PetscCheck(s3a == s3, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Right scaling vec non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT ".", s3a, s3);
10049566063dSJacob Faibussowitsch PetscCall(VecGetArrayRead(rr, &ar));
10056f08ad05SPierre Jolivet if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
10069566063dSJacob Faibussowitsch PetscCall(VecGetArray(mdn->lvec, &r));
10079566063dSJacob Faibussowitsch PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, ar, r, MPI_REPLACE));
10089566063dSJacob Faibussowitsch PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ar, r, MPI_REPLACE));
10099566063dSJacob Faibussowitsch PetscCall(VecRestoreArrayRead(rr, &ar));
10105b2fa520SLois Curfman McInnes for (i = 0; i < n; i++) {
10115b2fa520SLois Curfman McInnes x = r[i];
1012637a0070SStefano Zampini v = vv + i * lda;
10132205254eSKarl Rupp for (j = 0; j < m; j++) (*v++) *= x;
10145b2fa520SLois Curfman McInnes }
10159566063dSJacob Faibussowitsch PetscCall(VecRestoreArray(mdn->lvec, &r));
10169566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(1.0 * n * m));
10175b2fa520SLois Curfman McInnes }
10189566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(mdn->A, &vv));
10193ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
10205b2fa520SLois Curfman McInnes }
10215b2fa520SLois Curfman McInnes
MatNorm_MPIDense(Mat A,NormType type,PetscReal * nrm)102266976f2fSJacob Faibussowitsch static PetscErrorCode MatNorm_MPIDense(Mat A, NormType type, PetscReal *nrm)
1023d71ae5a4SJacob Faibussowitsch {
10243501a2bdSLois Curfman McInnes Mat_MPIDense *mdn = (Mat_MPIDense *)A->data;
102513f74950SBarry Smith PetscInt i, j;
1026e91c04dfSPierre Jolivet PetscMPIInt size;
1027329f5518SBarry Smith PetscReal sum = 0.0;
1028637a0070SStefano Zampini const PetscScalar *av, *v;
10293501a2bdSLois Curfman McInnes
10303a40ed3dSBarry Smith PetscFunctionBegin;
10319566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(mdn->A, &av));
1032637a0070SStefano Zampini v = av;
10339566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
1034616b8fbbSStefano Zampini if (size == 1) {
10359566063dSJacob Faibussowitsch PetscCall(MatNorm(mdn->A, type, nrm));
10363501a2bdSLois Curfman McInnes } else {
10373501a2bdSLois Curfman McInnes if (type == NORM_FROBENIUS) {
1038d0f46423SBarry Smith for (i = 0; i < mdn->A->cmap->n * mdn->A->rmap->n; i++) {
10399371c9d4SSatish Balay sum += PetscRealPart(PetscConj(*v) * (*v));
10409371c9d4SSatish Balay v++;
10413501a2bdSLois Curfman McInnes }
1042462c564dSBarry Smith PetscCallMPI(MPIU_Allreduce(&sum, nrm, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
10438f1a2a5eSBarry Smith *nrm = PetscSqrtReal(*nrm);
10449566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0 * mdn->A->cmap->n * mdn->A->rmap->n));
10453a40ed3dSBarry Smith } else if (type == NORM_1) {
1046e91c04dfSPierre Jolivet PetscReal *tmp;
1047e91c04dfSPierre Jolivet
1048e91c04dfSPierre Jolivet PetscCall(PetscCalloc1(A->cmap->N, &tmp));
1049064f8208SBarry Smith *nrm = 0.0;
1050637a0070SStefano Zampini v = av;
1051d0f46423SBarry Smith for (j = 0; j < mdn->A->cmap->n; j++) {
1052d0f46423SBarry Smith for (i = 0; i < mdn->A->rmap->n; i++) {
10539371c9d4SSatish Balay tmp[j] += PetscAbsScalar(*v);
10549371c9d4SSatish Balay v++;
10553501a2bdSLois Curfman McInnes }
10563501a2bdSLois Curfman McInnes }
1057e91c04dfSPierre Jolivet PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, tmp, A->cmap->N, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
1058d0f46423SBarry Smith for (j = 0; j < A->cmap->N; j++) {
1059e91c04dfSPierre Jolivet if (tmp[j] > *nrm) *nrm = tmp[j];
10603501a2bdSLois Curfman McInnes }
1061e91c04dfSPierre Jolivet PetscCall(PetscFree(tmp));
10629566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(A->cmap->n * A->rmap->n));
10633a40ed3dSBarry Smith } else if (type == NORM_INFINITY) { /* max row norm */
1064e91c04dfSPierre Jolivet PetscCall(MatNorm(mdn->A, type, nrm));
1065e91c04dfSPierre Jolivet PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, nrm, 1, MPIU_REAL, MPIU_MAX, PetscObjectComm((PetscObject)A)));
1066ce94432eSBarry Smith } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "No support for two norm");
10673501a2bdSLois Curfman McInnes }
10689566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(mdn->A, &av));
10693ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
10703501a2bdSLois Curfman McInnes }
10713501a2bdSLois Curfman McInnes
MatTranspose_MPIDense(Mat A,MatReuse reuse,Mat * matout)107266976f2fSJacob Faibussowitsch static PetscErrorCode MatTranspose_MPIDense(Mat A, MatReuse reuse, Mat *matout)
1073d71ae5a4SJacob Faibussowitsch {
10743501a2bdSLois Curfman McInnes Mat_MPIDense *a = (Mat_MPIDense *)A->data;
10753501a2bdSLois Curfman McInnes Mat B;
1076d0f46423SBarry Smith PetscInt M = A->rmap->N, N = A->cmap->N, m, n, *rwork, rstart = A->rmap->rstart;
1077637a0070SStefano Zampini PetscInt j, i, lda;
107887828ca2SBarry Smith PetscScalar *v;
10793501a2bdSLois Curfman McInnes
10803a40ed3dSBarry Smith PetscFunctionBegin;
10817fb60732SBarry Smith if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(A, *matout));
1082cf37664fSBarry Smith if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX) {
10839566063dSJacob Faibussowitsch PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
10849566063dSJacob Faibussowitsch PetscCall(MatSetSizes(B, A->cmap->n, A->rmap->n, N, M));
10859566063dSJacob Faibussowitsch PetscCall(MatSetType(B, ((PetscObject)A)->type_name));
10869566063dSJacob Faibussowitsch PetscCall(MatMPIDenseSetPreallocation(B, NULL));
1087637a0070SStefano Zampini } else B = *matout;
10883501a2bdSLois Curfman McInnes
10899371c9d4SSatish Balay m = a->A->rmap->n;
10909371c9d4SSatish Balay n = a->A->cmap->n;
10919566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)&v));
10929566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &lda));
10939566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(m, &rwork));
10943501a2bdSLois Curfman McInnes for (i = 0; i < m; i++) rwork[i] = rstart + i;
10951acff37aSSatish Balay for (j = 0; j < n; j++) {
10969566063dSJacob Faibussowitsch PetscCall(MatSetValues(B, 1, &j, m, rwork, v, INSERT_VALUES));
10978e3a54c0SPierre Jolivet v = PetscSafePointerPlusOffset(v, lda);
10983501a2bdSLois Curfman McInnes }
10999566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)&v));
11009566063dSJacob Faibussowitsch PetscCall(PetscFree(rwork));
11019566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
11029566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
1103cf37664fSBarry Smith if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_REUSE_MATRIX) {
11043501a2bdSLois Curfman McInnes *matout = B;
11053501a2bdSLois Curfman McInnes } else {
11069566063dSJacob Faibussowitsch PetscCall(MatHeaderMerge(A, &B));
11073501a2bdSLois Curfman McInnes }
11083ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1109096963f5SLois Curfman McInnes }
1110096963f5SLois Curfman McInnes
11116849ba73SBarry Smith static PetscErrorCode MatDuplicate_MPIDense(Mat, MatDuplicateOption, Mat *);
111252c5f739Sprj- PETSC_INTERN PetscErrorCode MatScale_MPIDense(Mat, PetscScalar);
11138965ea79SLois Curfman McInnes
MatSetUp_MPIDense(Mat A)111466976f2fSJacob Faibussowitsch static PetscErrorCode MatSetUp_MPIDense(Mat A)
1115d71ae5a4SJacob Faibussowitsch {
1116273d9f13SBarry Smith PetscFunctionBegin;
11179566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(A->rmap));
11189566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(A->cmap));
111948a46eb9SPierre Jolivet if (!A->preallocated) PetscCall(MatMPIDenseSetPreallocation(A, NULL));
11203ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1121273d9f13SBarry Smith }
1122273d9f13SBarry Smith
MatAXPY_MPIDense(Mat Y,PetscScalar alpha,Mat X,MatStructure str)112366976f2fSJacob Faibussowitsch static PetscErrorCode MatAXPY_MPIDense(Mat Y, PetscScalar alpha, Mat X, MatStructure str)
1124d71ae5a4SJacob Faibussowitsch {
1125488007eeSBarry Smith Mat_MPIDense *A = (Mat_MPIDense *)Y->data, *B = (Mat_MPIDense *)X->data;
1126488007eeSBarry Smith
1127488007eeSBarry Smith PetscFunctionBegin;
11289566063dSJacob Faibussowitsch PetscCall(MatAXPY(A->A, alpha, B->A, str));
11293ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1130488007eeSBarry Smith }
1131488007eeSBarry Smith
MatConjugate_MPIDense(Mat mat)113266976f2fSJacob Faibussowitsch static PetscErrorCode MatConjugate_MPIDense(Mat mat)
1133d71ae5a4SJacob Faibussowitsch {
1134ba337c44SJed Brown Mat_MPIDense *a = (Mat_MPIDense *)mat->data;
1135ba337c44SJed Brown
1136ba337c44SJed Brown PetscFunctionBegin;
11379566063dSJacob Faibussowitsch PetscCall(MatConjugate(a->A));
11383ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1139ba337c44SJed Brown }
1140ba337c44SJed Brown
MatRealPart_MPIDense(Mat A)114166976f2fSJacob Faibussowitsch static PetscErrorCode MatRealPart_MPIDense(Mat A)
1142d71ae5a4SJacob Faibussowitsch {
1143ba337c44SJed Brown Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1144ba337c44SJed Brown
1145ba337c44SJed Brown PetscFunctionBegin;
11469566063dSJacob Faibussowitsch PetscCall(MatRealPart(a->A));
11473ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1148ba337c44SJed Brown }
1149ba337c44SJed Brown
MatImaginaryPart_MPIDense(Mat A)115066976f2fSJacob Faibussowitsch static PetscErrorCode MatImaginaryPart_MPIDense(Mat A)
1151d71ae5a4SJacob Faibussowitsch {
1152ba337c44SJed Brown Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1153ba337c44SJed Brown
1154ba337c44SJed Brown PetscFunctionBegin;
11559566063dSJacob Faibussowitsch PetscCall(MatImaginaryPart(a->A));
11563ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1157ba337c44SJed Brown }
1158ba337c44SJed Brown
MatGetColumnVector_MPIDense(Mat A,Vec v,PetscInt col)1159d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatGetColumnVector_MPIDense(Mat A, Vec v, PetscInt col)
1160d71ae5a4SJacob Faibussowitsch {
1161637a0070SStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data;
116249a6ff4bSBarry Smith
116349a6ff4bSBarry Smith PetscFunctionBegin;
116428b400f6SJacob Faibussowitsch PetscCheck(a->A, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing local matrix");
116528b400f6SJacob Faibussowitsch PetscCheck(a->A->ops->getcolumnvector, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing get column operation");
11669566063dSJacob Faibussowitsch PetscCall((*a->A->ops->getcolumnvector)(a->A, v, col));
11673ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
116849a6ff4bSBarry Smith }
116949a6ff4bSBarry Smith
1170857cbf51SRichard Tran Mills PETSC_INTERN PetscErrorCode MatGetColumnReductions_SeqDense(Mat, PetscInt, PetscReal *);
117152c5f739Sprj-
MatGetColumnReductions_MPIDense(Mat A,PetscInt type,PetscReal * reductions)117266976f2fSJacob Faibussowitsch static PetscErrorCode MatGetColumnReductions_MPIDense(Mat A, PetscInt type, PetscReal *reductions)
1173d71ae5a4SJacob Faibussowitsch {
1174a873a8cdSSam Reynolds PetscInt i, m, n;
11750716a85fSBarry Smith Mat_MPIDense *a = (Mat_MPIDense *)A->data;
11760716a85fSBarry Smith
11770716a85fSBarry Smith PetscFunctionBegin;
11789566063dSJacob Faibussowitsch PetscCall(MatGetSize(A, &m, &n));
1179857cbf51SRichard Tran Mills if (type == REDUCTION_MEAN_REALPART) {
1180e91c04dfSPierre Jolivet PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_REALPART, reductions));
1181857cbf51SRichard Tran Mills } else if (type == REDUCTION_MEAN_IMAGINARYPART) {
1182e91c04dfSPierre Jolivet PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_IMAGINARYPART, reductions));
1183a873a8cdSSam Reynolds } else {
1184e91c04dfSPierre Jolivet PetscCall(MatGetColumnReductions_SeqDense(a->A, type, reductions));
1185a873a8cdSSam Reynolds }
1186857cbf51SRichard Tran Mills if (type == NORM_2) {
1187e91c04dfSPierre Jolivet for (i = 0; i < n; i++) reductions[i] *= reductions[i];
11880716a85fSBarry Smith }
1189857cbf51SRichard Tran Mills if (type == NORM_INFINITY) {
1190e91c04dfSPierre Jolivet PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, reductions, n, MPIU_REAL, MPIU_MAX, A->hdr.comm));
11910716a85fSBarry Smith } else {
1192e91c04dfSPierre Jolivet PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, reductions, n, MPIU_REAL, MPIU_SUM, A->hdr.comm));
11930716a85fSBarry Smith }
1194857cbf51SRichard Tran Mills if (type == NORM_2) {
1195a873a8cdSSam Reynolds for (i = 0; i < n; i++) reductions[i] = PetscSqrtReal(reductions[i]);
1196857cbf51SRichard Tran Mills } else if (type == REDUCTION_MEAN_REALPART || type == REDUCTION_MEAN_IMAGINARYPART) {
1197a873a8cdSSam Reynolds for (i = 0; i < n; i++) reductions[i] /= m;
11980716a85fSBarry Smith }
11993ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
12000716a85fSBarry Smith }
12010716a85fSBarry Smith
MatSetRandom_MPIDense(Mat x,PetscRandom rctx)1202d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatSetRandom_MPIDense(Mat x, PetscRandom rctx)
1203d71ae5a4SJacob Faibussowitsch {
120473a71a0fSBarry Smith Mat_MPIDense *d = (Mat_MPIDense *)x->data;
120573a71a0fSBarry Smith
120673a71a0fSBarry Smith PetscFunctionBegin;
12079566063dSJacob Faibussowitsch PetscCall(MatSetRandom(d->A, rctx));
12083faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE)
12093faff063SStefano Zampini x->offloadmask = d->A->offloadmask;
12103faff063SStefano Zampini #endif
12113ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
121273a71a0fSBarry Smith }
121373a71a0fSBarry Smith
12144222ddf1SHong Zhang static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat);
1215cc48ffa7SToby Isaac static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat);
12166718818eSStefano Zampini static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat);
1217a1f9a5e6SStefano Zampini static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat);
12186718818eSStefano Zampini static PetscErrorCode MatEqual_MPIDense(Mat, Mat, PetscBool *);
12196718818eSStefano Zampini static PetscErrorCode MatLoad_MPIDense(Mat, PetscViewer);
12205d8c7819SPierre Jolivet static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat);
1221cc48ffa7SToby Isaac
122209dc0095SBarry Smith static struct _MatOps MatOps_Values = {MatSetValues_MPIDense,
122309dc0095SBarry Smith MatGetRow_MPIDense,
122409dc0095SBarry Smith MatRestoreRow_MPIDense,
122509dc0095SBarry Smith MatMult_MPIDense,
122697304618SKris Buschelman /* 4*/ MatMultAdd_MPIDense,
12277c922b88SBarry Smith MatMultTranspose_MPIDense,
12287c922b88SBarry Smith MatMultTransposeAdd_MPIDense,
1229f4259b30SLisandro Dalcin NULL,
1230f4259b30SLisandro Dalcin NULL,
1231f4259b30SLisandro Dalcin NULL,
1232f4259b30SLisandro Dalcin /* 10*/ NULL,
1233f4259b30SLisandro Dalcin NULL,
1234f4259b30SLisandro Dalcin NULL,
1235f4259b30SLisandro Dalcin NULL,
123609dc0095SBarry Smith MatTranspose_MPIDense,
123797304618SKris Buschelman /* 15*/ MatGetInfo_MPIDense,
12386e4ee0c6SHong Zhang MatEqual_MPIDense,
123909dc0095SBarry Smith MatGetDiagonal_MPIDense,
12405b2fa520SLois Curfman McInnes MatDiagonalScale_MPIDense,
124109dc0095SBarry Smith MatNorm_MPIDense,
124297304618SKris Buschelman /* 20*/ MatAssemblyBegin_MPIDense,
124309dc0095SBarry Smith MatAssemblyEnd_MPIDense,
124409dc0095SBarry Smith MatSetOption_MPIDense,
124509dc0095SBarry Smith MatZeroEntries_MPIDense,
1246d519adbfSMatthew Knepley /* 24*/ MatZeroRows_MPIDense,
1247f4259b30SLisandro Dalcin NULL,
1248f4259b30SLisandro Dalcin NULL,
1249f4259b30SLisandro Dalcin NULL,
1250f4259b30SLisandro Dalcin NULL,
12514994cf47SJed Brown /* 29*/ MatSetUp_MPIDense,
1252f4259b30SLisandro Dalcin NULL,
1253f4259b30SLisandro Dalcin NULL,
1254c56a70eeSBarry Smith MatGetDiagonalBlock_MPIDense,
1255f4259b30SLisandro Dalcin NULL,
1256d519adbfSMatthew Knepley /* 34*/ MatDuplicate_MPIDense,
1257f4259b30SLisandro Dalcin NULL,
1258f4259b30SLisandro Dalcin NULL,
1259f4259b30SLisandro Dalcin NULL,
1260f4259b30SLisandro Dalcin NULL,
1261d519adbfSMatthew Knepley /* 39*/ MatAXPY_MPIDense,
12627dae84e0SHong Zhang MatCreateSubMatrices_MPIDense,
1263f4259b30SLisandro Dalcin NULL,
126409dc0095SBarry Smith MatGetValues_MPIDense,
1265a76f77c3SStefano Zampini MatCopy_MPIDense,
1266f4259b30SLisandro Dalcin /* 44*/ NULL,
126709dc0095SBarry Smith MatScale_MPIDense,
12682f605a99SJose E. Roman MatShift_MPIDense,
1269f4259b30SLisandro Dalcin NULL,
1270f4259b30SLisandro Dalcin NULL,
127173a71a0fSBarry Smith /* 49*/ MatSetRandom_MPIDense,
1272f4259b30SLisandro Dalcin NULL,
1273f4259b30SLisandro Dalcin NULL,
1274f4259b30SLisandro Dalcin NULL,
1275f4259b30SLisandro Dalcin NULL,
1276f4259b30SLisandro Dalcin /* 54*/ NULL,
1277f4259b30SLisandro Dalcin NULL,
1278f4259b30SLisandro Dalcin NULL,
1279f4259b30SLisandro Dalcin NULL,
1280f4259b30SLisandro Dalcin NULL,
12817dae84e0SHong Zhang /* 59*/ MatCreateSubMatrix_MPIDense,
1282b9b97703SBarry Smith MatDestroy_MPIDense,
1283b9b97703SBarry Smith MatView_MPIDense,
1284f4259b30SLisandro Dalcin NULL,
1285f4259b30SLisandro Dalcin NULL,
1286f4259b30SLisandro Dalcin /* 64*/ NULL,
1287f4259b30SLisandro Dalcin NULL,
1288f4259b30SLisandro Dalcin NULL,
1289f4259b30SLisandro Dalcin NULL,
1290f4259b30SLisandro Dalcin NULL,
1291f4259b30SLisandro Dalcin /* 69*/ NULL,
1292f4259b30SLisandro Dalcin NULL,
1293f4259b30SLisandro Dalcin NULL,
1294f4259b30SLisandro Dalcin NULL,
1295f4259b30SLisandro Dalcin NULL,
1296f4259b30SLisandro Dalcin /* 74*/ NULL,
1297f4259b30SLisandro Dalcin NULL,
1298f4259b30SLisandro Dalcin NULL,
1299f4259b30SLisandro Dalcin NULL,
13008bb0f5c6SPierre Jolivet MatLoad_MPIDense,
1301f4259b30SLisandro Dalcin /* 79*/ NULL,
1302f4259b30SLisandro Dalcin NULL,
1303f4259b30SLisandro Dalcin NULL,
1304f4259b30SLisandro Dalcin NULL,
13058bb0f5c6SPierre Jolivet /* 83*/ NULL,
1306f4259b30SLisandro Dalcin NULL,
1307f4259b30SLisandro Dalcin NULL,
1308f4259b30SLisandro Dalcin NULL,
13096718818eSStefano Zampini MatMatTransposeMultSymbolic_MPIDense_MPIDense,
1310cc48ffa7SToby Isaac MatMatTransposeMultNumeric_MPIDense_MPIDense,
13118bb0f5c6SPierre Jolivet /* 89*/ NULL,
13128bb0f5c6SPierre Jolivet MatProductSetFromOptions_MPIDense,
1313f4259b30SLisandro Dalcin NULL,
1314f4259b30SLisandro Dalcin NULL,
1315ba337c44SJed Brown MatConjugate_MPIDense,
13168bb0f5c6SPierre Jolivet /* 94*/ NULL,
1317f4259b30SLisandro Dalcin NULL,
1318ba337c44SJed Brown MatRealPart_MPIDense,
1319ba337c44SJed Brown MatImaginaryPart_MPIDense,
1320f4259b30SLisandro Dalcin NULL,
13218bb0f5c6SPierre Jolivet /*99*/ NULL,
13228bb0f5c6SPierre Jolivet NULL,
13238bb0f5c6SPierre Jolivet NULL,
13248bb0f5c6SPierre Jolivet NULL,
13258bb0f5c6SPierre Jolivet MatGetColumnVector_MPIDense,
1326421480d9SBarry Smith /*104*/ NULL,
13278bb0f5c6SPierre Jolivet NULL,
13288bb0f5c6SPierre Jolivet NULL,
13298bb0f5c6SPierre Jolivet NULL,
1330f4259b30SLisandro Dalcin NULL,
1331f4259b30SLisandro Dalcin /*109*/ NULL,
1332f4259b30SLisandro Dalcin NULL,
1333459e8d23SBlanca Mellado Pinto MatMultHermitianTranspose_MPIDense,
1334459e8d23SBlanca Mellado Pinto MatMultHermitianTransposeAdd_MPIDense,
13358bb0f5c6SPierre Jolivet NULL,
1336421480d9SBarry Smith /*114*/ NULL,
13378bb0f5c6SPierre Jolivet MatGetColumnReductions_MPIDense,
13388bb0f5c6SPierre Jolivet NULL,
13398bb0f5c6SPierre Jolivet NULL,
13408bb0f5c6SPierre Jolivet NULL,
13418bb0f5c6SPierre Jolivet /*120*/ NULL,
13428bb0f5c6SPierre Jolivet MatTransposeMatMultSymbolic_MPIDense_MPIDense,
13438bb0f5c6SPierre Jolivet MatTransposeMatMultNumeric_MPIDense_MPIDense,
1344f4259b30SLisandro Dalcin NULL,
1345f4259b30SLisandro Dalcin /*124*/ NULL,
13468bb0f5c6SPierre Jolivet NULL,
1347f4259b30SLisandro Dalcin NULL,
1348f4259b30SLisandro Dalcin NULL,
1349f4259b30SLisandro Dalcin NULL,
1350f4259b30SLisandro Dalcin /*129*/ NULL,
1351f4259b30SLisandro Dalcin NULL,
13528bb0f5c6SPierre Jolivet MatCreateMPIMatConcatenateSeqMat_MPIDense,
13538bb0f5c6SPierre Jolivet NULL,
1354f4259b30SLisandro Dalcin NULL,
1355f4259b30SLisandro Dalcin /*134*/ NULL,
1356f4259b30SLisandro Dalcin NULL,
1357f4259b30SLisandro Dalcin NULL,
1358f4259b30SLisandro Dalcin NULL,
1359f4259b30SLisandro Dalcin NULL,
1360f4259b30SLisandro Dalcin /*139*/ NULL,
1361f4259b30SLisandro Dalcin NULL,
1362f4259b30SLisandro Dalcin NULL,
136303db1824SAlex Lindsay NULL,
1364c2be7ffeSStefano Zampini NULL,
1365dec0b466SHong Zhang NULL};
13668965ea79SLois Curfman McInnes
MatMPIDenseSetPreallocation_MPIDense(Mat mat,PetscScalar * data)136766976f2fSJacob Faibussowitsch static PetscErrorCode MatMPIDenseSetPreallocation_MPIDense(Mat mat, PetscScalar *data)
1368d71ae5a4SJacob Faibussowitsch {
1369637a0070SStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)mat->data;
137047d993e7Ssuyashtn MatType mtype = MATSEQDENSE;
1371a23d5eceSKris Buschelman
1372a23d5eceSKris Buschelman PetscFunctionBegin;
137328b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
13749566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(mat->rmap));
13759566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(mat->cmap));
1376637a0070SStefano Zampini if (!a->A) {
13779566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
13789566063dSJacob Faibussowitsch PetscCall(MatSetSizes(a->A, mat->rmap->n, mat->cmap->N, mat->rmap->n, mat->cmap->N));
1379637a0070SStefano Zampini }
138047d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA)
138147d993e7Ssuyashtn PetscBool iscuda;
13829566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSECUDA, &iscuda));
138347d993e7Ssuyashtn if (iscuda) mtype = MATSEQDENSECUDA;
138447d993e7Ssuyashtn #endif
138547d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
138647d993e7Ssuyashtn PetscBool iship;
138747d993e7Ssuyashtn PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSEHIP, &iship));
138847d993e7Ssuyashtn if (iship) mtype = MATSEQDENSEHIP;
138947d993e7Ssuyashtn #endif
139047d993e7Ssuyashtn PetscCall(MatSetType(a->A, mtype));
13919566063dSJacob Faibussowitsch PetscCall(MatSeqDenseSetPreallocation(a->A, data));
139247d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
13933faff063SStefano Zampini mat->offloadmask = a->A->offloadmask;
13943faff063SStefano Zampini #endif
1395637a0070SStefano Zampini mat->preallocated = PETSC_TRUE;
13966f08ad05SPierre Jolivet mat->assembled = PETSC_TRUE;
13973ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1398a23d5eceSKris Buschelman }
1399a23d5eceSKris Buschelman
MatConvert_MPIAIJ_MPIDense(Mat A,MatType newtype,MatReuse reuse,Mat * newmat)1400d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode MatConvert_MPIAIJ_MPIDense(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1401d71ae5a4SJacob Faibussowitsch {
14024e29119aSPierre Jolivet Mat B, C;
14034e29119aSPierre Jolivet
14044e29119aSPierre Jolivet PetscFunctionBegin;
14059566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A, MAT_INITIAL_MATRIX, &C));
14069566063dSJacob Faibussowitsch PetscCall(MatConvert_SeqAIJ_SeqDense(C, MATSEQDENSE, MAT_INITIAL_MATRIX, &B));
14079566063dSJacob Faibussowitsch PetscCall(MatDestroy(&C));
14084e29119aSPierre Jolivet if (reuse == MAT_REUSE_MATRIX) {
14094e29119aSPierre Jolivet C = *newmat;
14104e29119aSPierre Jolivet } else C = NULL;
14119566063dSJacob Faibussowitsch PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C));
14129566063dSJacob Faibussowitsch PetscCall(MatDestroy(&B));
1413ac530a7eSPierre Jolivet if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &C));
1414ac530a7eSPierre Jolivet else if (reuse == MAT_INITIAL_MATRIX) *newmat = C;
14153ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
14164e29119aSPierre Jolivet }
14174e29119aSPierre Jolivet
MatConvert_MPIDense_MPIAIJ(Mat A,MatType newtype,MatReuse reuse,Mat * newmat)141866976f2fSJacob Faibussowitsch static PetscErrorCode MatConvert_MPIDense_MPIAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1419d71ae5a4SJacob Faibussowitsch {
14204e29119aSPierre Jolivet Mat B, C;
14214e29119aSPierre Jolivet
14224e29119aSPierre Jolivet PetscFunctionBegin;
14239566063dSJacob Faibussowitsch PetscCall(MatDenseGetLocalMatrix(A, &C));
14249566063dSJacob Faibussowitsch PetscCall(MatConvert_SeqDense_SeqAIJ(C, MATSEQAIJ, MAT_INITIAL_MATRIX, &B));
14254e29119aSPierre Jolivet if (reuse == MAT_REUSE_MATRIX) {
14264e29119aSPierre Jolivet C = *newmat;
14274e29119aSPierre Jolivet } else C = NULL;
14289566063dSJacob Faibussowitsch PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C));
14299566063dSJacob Faibussowitsch PetscCall(MatDestroy(&B));
1430ac530a7eSPierre Jolivet if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &C));
1431ac530a7eSPierre Jolivet else if (reuse == MAT_INITIAL_MATRIX) *newmat = C;
14323ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
14334e29119aSPierre Jolivet }
14344e29119aSPierre Jolivet
143565b80a83SHong Zhang #if defined(PETSC_HAVE_ELEMENTAL)
MatConvert_MPIDense_Elemental(Mat A,MatType newtype,MatReuse reuse,Mat * newmat)1436d71ae5a4SJacob Faibussowitsch PETSC_INTERN PetscErrorCode MatConvert_MPIDense_Elemental(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1437d71ae5a4SJacob Faibussowitsch {
1438f709e8d7SPierre Jolivet Mat_MPIDense *a = (Mat_MPIDense *)A->data;
14398ea901baSHong Zhang Mat mat_elemental;
144032d7a744SHong Zhang PetscScalar *v;
1441f709e8d7SPierre Jolivet PetscInt m = A->rmap->n, N = A->cmap->N, rstart = A->rmap->rstart, i, *rows, *cols, lda;
14428ea901baSHong Zhang
14438baccfbdSHong Zhang PetscFunctionBegin;
1444378336b6SHong Zhang if (reuse == MAT_REUSE_MATRIX) {
1445378336b6SHong Zhang mat_elemental = *newmat;
14469566063dSJacob Faibussowitsch PetscCall(MatZeroEntries(*newmat));
1447378336b6SHong Zhang } else {
14489566063dSJacob Faibussowitsch PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat_elemental));
14499566063dSJacob Faibussowitsch PetscCall(MatSetSizes(mat_elemental, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N));
14509566063dSJacob Faibussowitsch PetscCall(MatSetType(mat_elemental, MATELEMENTAL));
14519566063dSJacob Faibussowitsch PetscCall(MatSetUp(mat_elemental));
14529566063dSJacob Faibussowitsch PetscCall(MatSetOption(mat_elemental, MAT_ROW_ORIENTED, PETSC_FALSE));
1453378336b6SHong Zhang }
1454378336b6SHong Zhang
14559566063dSJacob Faibussowitsch PetscCall(PetscMalloc2(m, &rows, N, &cols));
145632d7a744SHong Zhang for (i = 0; i < N; i++) cols[i] = i;
145732d7a744SHong Zhang for (i = 0; i < m; i++) rows[i] = rstart + i;
14588ea901baSHong Zhang
1459637a0070SStefano Zampini /* PETSc-Elemental interface uses axpy for setting off-processor entries, only ADD_VALUES is allowed */
14609566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(A, &v));
1461f709e8d7SPierre Jolivet PetscCall(MatDenseGetLDA(a->A, &lda));
1462f709e8d7SPierre Jolivet if (lda == m) PetscCall(MatSetValues(mat_elemental, m, rows, N, cols, v, ADD_VALUES));
1463f709e8d7SPierre Jolivet else {
1464f709e8d7SPierre Jolivet for (i = 0; i < N; i++) PetscCall(MatSetValues(mat_elemental, m, rows, 1, &i, v + lda * i, ADD_VALUES));
1465f709e8d7SPierre Jolivet }
14669566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(mat_elemental, MAT_FINAL_ASSEMBLY));
14679566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(mat_elemental, MAT_FINAL_ASSEMBLY));
14689566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(A, &v));
14699566063dSJacob Faibussowitsch PetscCall(PetscFree2(rows, cols));
14708ea901baSHong Zhang
1471511c6705SHong Zhang if (reuse == MAT_INPLACE_MATRIX) {
14729566063dSJacob Faibussowitsch PetscCall(MatHeaderReplace(A, &mat_elemental));
14738ea901baSHong Zhang } else {
14748ea901baSHong Zhang *newmat = mat_elemental;
14758ea901baSHong Zhang }
14763ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
14778baccfbdSHong Zhang }
147865b80a83SHong Zhang #endif
14798baccfbdSHong Zhang
MatDenseGetColumn_MPIDense(Mat A,PetscInt col,PetscScalar ** vals)1480d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseGetColumn_MPIDense(Mat A, PetscInt col, PetscScalar **vals)
1481d71ae5a4SJacob Faibussowitsch {
148286aefd0dSHong Zhang Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
148386aefd0dSHong Zhang
148486aefd0dSHong Zhang PetscFunctionBegin;
14859566063dSJacob Faibussowitsch PetscCall(MatDenseGetColumn(mat->A, col, vals));
14863ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
148786aefd0dSHong Zhang }
148886aefd0dSHong Zhang
MatDenseRestoreColumn_MPIDense(Mat A,PetscScalar ** vals)1489d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDenseRestoreColumn_MPIDense(Mat A, PetscScalar **vals)
1490d71ae5a4SJacob Faibussowitsch {
149186aefd0dSHong Zhang Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
149286aefd0dSHong Zhang
149386aefd0dSHong Zhang PetscFunctionBegin;
14949566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreColumn(mat->A, vals));
14953ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
149686aefd0dSHong Zhang }
149786aefd0dSHong Zhang
MatCreateMPIMatConcatenateSeqMat_MPIDense(MPI_Comm comm,Mat inmat,PetscInt n,MatReuse scall,Mat * outmat)1498d71ae5a4SJacob Faibussowitsch PetscErrorCode MatCreateMPIMatConcatenateSeqMat_MPIDense(MPI_Comm comm, Mat inmat, PetscInt n, MatReuse scall, Mat *outmat)
1499d71ae5a4SJacob Faibussowitsch {
150094e2cb23SJakub Kruzik Mat_MPIDense *mat;
150194e2cb23SJakub Kruzik PetscInt m, nloc, N;
150294e2cb23SJakub Kruzik
150394e2cb23SJakub Kruzik PetscFunctionBegin;
15049566063dSJacob Faibussowitsch PetscCall(MatGetSize(inmat, &m, &N));
15059566063dSJacob Faibussowitsch PetscCall(MatGetLocalSize(inmat, NULL, &nloc));
150694e2cb23SJakub Kruzik if (scall == MAT_INITIAL_MATRIX) { /* symbolic phase */
150794e2cb23SJakub Kruzik PetscInt sum;
150894e2cb23SJakub Kruzik
150948a46eb9SPierre Jolivet if (n == PETSC_DECIDE) PetscCall(PetscSplitOwnership(comm, &n, &N));
151094e2cb23SJakub Kruzik /* Check sum(n) = N */
1511462c564dSBarry Smith PetscCallMPI(MPIU_Allreduce(&n, &sum, 1, MPIU_INT, MPI_SUM, comm));
151208401ef6SPierre Jolivet PetscCheck(sum == N, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Sum of local columns %" PetscInt_FMT " != global columns %" PetscInt_FMT, sum, N);
151394e2cb23SJakub Kruzik
15149566063dSJacob Faibussowitsch PetscCall(MatCreateDense(comm, m, n, PETSC_DETERMINE, N, NULL, outmat));
15159566063dSJacob Faibussowitsch PetscCall(MatSetOption(*outmat, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
151694e2cb23SJakub Kruzik }
151794e2cb23SJakub Kruzik
151894e2cb23SJakub Kruzik /* numeric phase */
151994e2cb23SJakub Kruzik mat = (Mat_MPIDense *)(*outmat)->data;
15209566063dSJacob Faibussowitsch PetscCall(MatCopy(inmat, mat->A, SAME_NONZERO_PATTERN));
15213ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
152294e2cb23SJakub Kruzik }
152394e2cb23SJakub Kruzik
MatDenseGetColumnVec_MPIDense(Mat A,PetscInt col,Vec * v)1524d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVec_MPIDense(Mat A, PetscInt col, Vec *v)
1525d71ae5a4SJacob Faibussowitsch {
15266947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15276947451fSStefano Zampini PetscInt lda;
15286947451fSStefano Zampini
15296947451fSStefano Zampini PetscFunctionBegin;
153028b400f6SJacob Faibussowitsch PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
153128b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1532d16ceb75SStefano Zampini if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
15336947451fSStefano Zampini a->vecinuse = col + 1;
15349566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &lda));
15359566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(a->A, (PetscScalar **)&a->ptrinuse));
15368e3a54c0SPierre Jolivet PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
15376947451fSStefano Zampini *v = a->cvec;
15383ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
15396947451fSStefano Zampini }
15406947451fSStefano Zampini
MatDenseRestoreColumnVec_MPIDense(Mat A,PetscInt col,Vec * v)1541d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVec_MPIDense(Mat A, PetscInt col, Vec *v)
1542d71ae5a4SJacob Faibussowitsch {
15436947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15446947451fSStefano Zampini
15456947451fSStefano Zampini PetscFunctionBegin;
154628b400f6SJacob Faibussowitsch PetscCheck(a->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
154728b400f6SJacob Faibussowitsch PetscCheck(a->cvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column vector");
15484186a4bbSPierre Jolivet VecCheckAssembled(a->cvec);
15496947451fSStefano Zampini a->vecinuse = 0;
15509566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(a->A, (PetscScalar **)&a->ptrinuse));
15519566063dSJacob Faibussowitsch PetscCall(VecResetArray(a->cvec));
1552742765d3SMatthew Knepley if (v) *v = NULL;
15533ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
15546947451fSStefano Zampini }
15556947451fSStefano Zampini
MatDenseGetColumnVecRead_MPIDense(Mat A,PetscInt col,Vec * v)1556d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v)
1557d71ae5a4SJacob Faibussowitsch {
15586947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15596947451fSStefano Zampini PetscInt lda;
15606947451fSStefano Zampini
15616947451fSStefano Zampini PetscFunctionBegin;
156228b400f6SJacob Faibussowitsch PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
156328b400f6SJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1564d16ceb75SStefano Zampini if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
15656947451fSStefano Zampini a->vecinuse = col + 1;
15669566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &lda));
15679566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(a->A, &a->ptrinuse));
15688e3a54c0SPierre Jolivet PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
15699566063dSJacob Faibussowitsch PetscCall(VecLockReadPush(a->cvec));
15706947451fSStefano Zampini *v = a->cvec;
15713ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
15726947451fSStefano Zampini }
15736947451fSStefano Zampini
MatDenseRestoreColumnVecRead_MPIDense(Mat A,PetscInt col,Vec * v)1574d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v)
1575d71ae5a4SJacob Faibussowitsch {
15766947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15776947451fSStefano Zampini
15786947451fSStefano Zampini PetscFunctionBegin;
15795f80ce2aSJacob Faibussowitsch PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
15805f80ce2aSJacob Faibussowitsch PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector");
15814186a4bbSPierre Jolivet VecCheckAssembled(a->cvec);
15826947451fSStefano Zampini a->vecinuse = 0;
15839566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(a->A, &a->ptrinuse));
15849566063dSJacob Faibussowitsch PetscCall(VecLockReadPop(a->cvec));
15859566063dSJacob Faibussowitsch PetscCall(VecResetArray(a->cvec));
1586742765d3SMatthew Knepley if (v) *v = NULL;
15873ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
15886947451fSStefano Zampini }
15896947451fSStefano Zampini
MatDenseGetColumnVecWrite_MPIDense(Mat A,PetscInt col,Vec * v)1590d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseGetColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v)
1591d71ae5a4SJacob Faibussowitsch {
15926947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data;
15936947451fSStefano Zampini PetscInt lda;
15946947451fSStefano Zampini
15956947451fSStefano Zampini PetscFunctionBegin;
15965f80ce2aSJacob Faibussowitsch PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
15975f80ce2aSJacob Faibussowitsch PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1598d16ceb75SStefano Zampini if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
15996947451fSStefano Zampini a->vecinuse = col + 1;
16009566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &lda));
16019566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayWrite(a->A, (PetscScalar **)&a->ptrinuse));
16028e3a54c0SPierre Jolivet PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
16036947451fSStefano Zampini *v = a->cvec;
16043ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
16056947451fSStefano Zampini }
16066947451fSStefano Zampini
MatDenseRestoreColumnVecWrite_MPIDense(Mat A,PetscInt col,Vec * v)1607d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseRestoreColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v)
1608d71ae5a4SJacob Faibussowitsch {
16096947451fSStefano Zampini Mat_MPIDense *a = (Mat_MPIDense *)A->data;
16106947451fSStefano Zampini
16116947451fSStefano Zampini PetscFunctionBegin;
16125f80ce2aSJacob Faibussowitsch PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
16135f80ce2aSJacob Faibussowitsch PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector");
16144186a4bbSPierre Jolivet VecCheckAssembled(a->cvec);
16156947451fSStefano Zampini a->vecinuse = 0;
16169566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayWrite(a->A, (PetscScalar **)&a->ptrinuse));
16179566063dSJacob Faibussowitsch PetscCall(VecResetArray(a->cvec));
1618742765d3SMatthew Knepley if (v) *v = NULL;
16193ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
16206947451fSStefano Zampini }
16216947451fSStefano Zampini
MatDenseGetSubMatrix_MPIDense(Mat A,PetscInt rbegin,PetscInt rend,PetscInt cbegin,PetscInt cend,Mat * v)162266976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseGetSubMatrix_MPIDense(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *v)
1623d71ae5a4SJacob Faibussowitsch {
16245ea7661aSPierre Jolivet Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1625616b8fbbSStefano Zampini Mat_MPIDense *c;
1626616b8fbbSStefano Zampini MPI_Comm comm;
1627d016bddeSToby Isaac PetscInt prbegin, prend, pcbegin, pcend;
16285ea7661aSPierre Jolivet
16295ea7661aSPierre Jolivet PetscFunctionBegin;
16309566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
16315f80ce2aSJacob Faibussowitsch PetscCheck(!a->vecinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
16325f80ce2aSJacob Faibussowitsch PetscCheck(!a->matinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1633d016bddeSToby Isaac prbegin = PetscMax(0, PetscMin(A->rmap->rend, rbegin) - A->rmap->rstart);
1634d016bddeSToby Isaac prend = PetscMin(A->rmap->n, PetscMax(0, rend - A->rmap->rstart));
1635d016bddeSToby Isaac pcbegin = PetscMax(0, PetscMin(A->cmap->rend, cbegin) - A->cmap->rstart);
1636d016bddeSToby Isaac pcend = PetscMin(A->cmap->n, PetscMax(0, cend - A->cmap->rstart));
16375ea7661aSPierre Jolivet if (!a->cmat) {
16389566063dSJacob Faibussowitsch PetscCall(MatCreate(comm, &a->cmat));
16399566063dSJacob Faibussowitsch PetscCall(MatSetType(a->cmat, ((PetscObject)A)->type_name));
1640a2748737SPierre Jolivet if (rend - rbegin == A->rmap->N) PetscCall(PetscLayoutReference(A->rmap, &a->cmat->rmap));
1641a2748737SPierre Jolivet else {
1642d016bddeSToby Isaac PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, prend - prbegin));
1643a2748737SPierre Jolivet PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin));
1644a2748737SPierre Jolivet PetscCall(PetscLayoutSetUp(a->cmat->rmap));
1645a2748737SPierre Jolivet }
1646d016bddeSToby Isaac if (cend - cbegin == A->cmap->N) PetscCall(PetscLayoutReference(A->cmap, &a->cmat->cmap));
1647d016bddeSToby Isaac else {
1648d016bddeSToby Isaac PetscCall(PetscLayoutSetLocalSize(a->cmat->cmap, pcend - pcbegin));
16499566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin));
16509566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(a->cmat->cmap));
1651a2748737SPierre Jolivet }
1652616b8fbbSStefano Zampini c = (Mat_MPIDense *)a->cmat->data;
1653d016bddeSToby Isaac c->sub_rbegin = rbegin;
1654d016bddeSToby Isaac c->sub_rend = rend;
1655d016bddeSToby Isaac c->sub_cbegin = cbegin;
1656d016bddeSToby Isaac c->sub_cend = cend;
1657d016bddeSToby Isaac }
1658d016bddeSToby Isaac c = (Mat_MPIDense *)a->cmat->data;
1659d016bddeSToby Isaac if (c->sub_rbegin != rbegin || c->sub_rend != rend) {
1660d016bddeSToby Isaac PetscCall(PetscLayoutDestroy(&a->cmat->rmap));
1661d016bddeSToby Isaac PetscCall(PetscLayoutCreate(comm, &a->cmat->rmap));
1662d016bddeSToby Isaac PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, prend - prbegin));
1663d016bddeSToby Isaac PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin));
1664d016bddeSToby Isaac PetscCall(PetscLayoutSetUp(a->cmat->rmap));
1665d016bddeSToby Isaac c->sub_rbegin = rbegin;
1666d016bddeSToby Isaac c->sub_rend = rend;
1667d016bddeSToby Isaac }
1668d016bddeSToby Isaac if (c->sub_cbegin != cbegin || c->sub_cend != cend) {
1669d016bddeSToby Isaac // special optimization: check if all columns are owned by rank 0, in which case no communication is necessary
1670d016bddeSToby Isaac if ((cend - cbegin != a->cmat->cmap->N) || (A->cmap->range[1] != A->cmap->N)) {
1671d016bddeSToby Isaac PetscCall(PetscLayoutDestroy(&a->cmat->cmap));
1672d016bddeSToby Isaac PetscCall(PetscLayoutCreate(comm, &a->cmat->cmap));
1673d016bddeSToby Isaac PetscCall(PetscLayoutSetLocalSize(a->cmat->cmap, pcend - pcbegin));
1674d016bddeSToby Isaac PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin));
1675d016bddeSToby Isaac PetscCall(PetscLayoutSetUp(a->cmat->cmap));
1676d016bddeSToby Isaac PetscCall(VecDestroy(&c->lvec));
1677d016bddeSToby Isaac PetscCall(PetscSFDestroy(&c->Mvctx));
1678d016bddeSToby Isaac }
1679d016bddeSToby Isaac c->sub_cbegin = cbegin;
1680d016bddeSToby Isaac c->sub_cend = cend;
1681d016bddeSToby Isaac }
16825f80ce2aSJacob Faibussowitsch PetscCheck(!c->A, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1683d016bddeSToby Isaac PetscCall(MatDenseGetSubMatrix(a->A, prbegin, prend, cbegin, cend, &c->A));
16843faff063SStefano Zampini
1685616b8fbbSStefano Zampini a->cmat->preallocated = PETSC_TRUE;
1686616b8fbbSStefano Zampini a->cmat->assembled = PETSC_TRUE;
16873faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE)
16883faff063SStefano Zampini a->cmat->offloadmask = c->A->offloadmask;
16893faff063SStefano Zampini #endif
16905ea7661aSPierre Jolivet a->matinuse = cbegin + 1;
16915ea7661aSPierre Jolivet *v = a->cmat;
16923ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
16935ea7661aSPierre Jolivet }
16945ea7661aSPierre Jolivet
MatDenseRestoreSubMatrix_MPIDense(Mat A,Mat * v)169566976f2fSJacob Faibussowitsch static PetscErrorCode MatDenseRestoreSubMatrix_MPIDense(Mat A, Mat *v)
1696d71ae5a4SJacob Faibussowitsch {
16975ea7661aSPierre Jolivet Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1698616b8fbbSStefano Zampini Mat_MPIDense *c;
16995ea7661aSPierre Jolivet
17005ea7661aSPierre Jolivet PetscFunctionBegin;
17015f80ce2aSJacob Faibussowitsch PetscCheck(a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetSubMatrix() first");
17025f80ce2aSJacob Faibussowitsch PetscCheck(a->cmat, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal matrix");
17035f80ce2aSJacob Faibussowitsch PetscCheck(*v == a->cmat, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Not the matrix obtained from MatDenseGetSubMatrix()");
17045ea7661aSPierre Jolivet a->matinuse = 0;
1705616b8fbbSStefano Zampini c = (Mat_MPIDense *)a->cmat->data;
17069566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreSubMatrix(a->A, &c->A));
1707beb1828fSPierre Jolivet *v = NULL;
17083faff063SStefano Zampini #if defined(PETSC_HAVE_DEVICE)
17093faff063SStefano Zampini A->offloadmask = a->A->offloadmask;
17103faff063SStefano Zampini #endif
17113ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
17125ea7661aSPierre Jolivet }
17135ea7661aSPierre Jolivet
1714002bc6daSRichard Tran Mills /*MC
1715002bc6daSRichard Tran Mills MATMPIDENSE - MATMPIDENSE = "mpidense" - A matrix type to be used for distributed dense matrices.
1716002bc6daSRichard Tran Mills
17172ef1f0ffSBarry Smith Options Database Key:
171811a5261eSBarry Smith . -mat_type mpidense - sets the matrix type to `MATMPIDENSE` during a call to `MatSetFromOptions()`
1719002bc6daSRichard Tran Mills
1720002bc6daSRichard Tran Mills Level: beginner
1721002bc6daSRichard Tran Mills
17221cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MatCreateDense()`, `MATSEQDENSE`, `MATDENSE`
1723002bc6daSRichard Tran Mills M*/
MatCreate_MPIDense(Mat mat)17244742e46bSJacob Faibussowitsch PetscErrorCode MatCreate_MPIDense(Mat mat)
1725d71ae5a4SJacob Faibussowitsch {
1726273d9f13SBarry Smith Mat_MPIDense *a;
1727273d9f13SBarry Smith
1728273d9f13SBarry Smith PetscFunctionBegin;
17294dfa11a4SJacob Faibussowitsch PetscCall(PetscNew(&a));
1730b0a32e0cSBarry Smith mat->data = (void *)a;
1731aea10558SJacob Faibussowitsch mat->ops[0] = MatOps_Values;
1732273d9f13SBarry Smith
1733273d9f13SBarry Smith mat->insertmode = NOT_SET_VALUES;
1734273d9f13SBarry Smith
1735273d9f13SBarry Smith /* build cache for off array entries formed */
1736273d9f13SBarry Smith a->donotstash = PETSC_FALSE;
17372205254eSKarl Rupp
17389566063dSJacob Faibussowitsch PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)mat), 1, &mat->stash));
1739273d9f13SBarry Smith
1740273d9f13SBarry Smith /* stuff used for matrix vector multiply */
1741f4259b30SLisandro Dalcin a->lvec = NULL;
1742f4259b30SLisandro Dalcin a->Mvctx = NULL;
1743273d9f13SBarry Smith a->roworiented = PETSC_TRUE;
1744273d9f13SBarry Smith
17459566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", MatDenseGetLDA_MPIDense));
17469566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", MatDenseSetLDA_MPIDense));
17479566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", MatDenseGetArray_MPIDense));
17489566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", MatDenseRestoreArray_MPIDense));
17499566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", MatDenseGetArrayRead_MPIDense));
17509566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", MatDenseRestoreArrayRead_MPIDense));
17519566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", MatDenseGetArrayWrite_MPIDense));
17529566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", MatDenseRestoreArrayWrite_MPIDense));
17539566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", MatDensePlaceArray_MPIDense));
17549566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", MatDenseResetArray_MPIDense));
17559566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", MatDenseReplaceArray_MPIDense));
17569566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_MPIDense));
17579566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_MPIDense));
17589566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_MPIDense));
17599566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_MPIDense));
17609566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_MPIDense));
17619566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_MPIDense));
17629566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_MPIDense));
17639566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_MPIDense));
17649566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", MatConvert_MPIAIJ_MPIDense));
17659566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", MatConvert_MPIDense_MPIAIJ));
17668baccfbdSHong Zhang #if defined(PETSC_HAVE_ELEMENTAL)
17679566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", MatConvert_MPIDense_Elemental));
17688baccfbdSHong Zhang #endif
1769d1a032dbSPierre Jolivet #if defined(PETSC_HAVE_SCALAPACK) && (defined(PETSC_USE_REAL_SINGLE) || defined(PETSC_USE_REAL_DOUBLE))
17709566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", MatConvert_Dense_ScaLAPACK));
1771d24d4204SJose E. Roman #endif
1772637a0070SStefano Zampini #if defined(PETSC_HAVE_CUDA)
17739566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", MatConvert_MPIDense_MPIDenseCUDA));
1774637a0070SStefano Zampini #endif
17759566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", MatMPIDenseSetPreallocation_MPIDense));
17769566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
17779566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
17786718818eSStefano Zampini #if defined(PETSC_HAVE_CUDA)
17799566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
17809566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
17816718818eSStefano Zampini #endif
178247d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
178347d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", MatConvert_MPIDense_MPIDenseHIP));
178447d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
178547d993e7Ssuyashtn PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
178647d993e7Ssuyashtn #endif
17879566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", MatDenseGetColumn_MPIDense));
17889566063dSJacob Faibussowitsch PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", MatDenseRestoreColumn_MPIDense));
1789d016bddeSToby Isaac PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultColumnRange_C", MatMultColumnRange_MPIDense));
17900be0d8bdSHansol Suh PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", MatMultAddColumnRange_MPIDense));
17910be0d8bdSHansol Suh PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", MatMultHermitianTransposeColumnRange_MPIDense));
17920be0d8bdSHansol Suh PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", MatMultHermitianTransposeAddColumnRange_MPIDense));
17939566063dSJacob Faibussowitsch PetscCall(PetscObjectChangeTypeName((PetscObject)mat, MATMPIDENSE));
17943ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1795273d9f13SBarry Smith }
1796273d9f13SBarry Smith
1797209238afSKris Buschelman /*MC
1798002d173eSKris Buschelman MATDENSE - MATDENSE = "dense" - A matrix type to be used for dense matrices.
1799209238afSKris Buschelman
180011a5261eSBarry Smith This matrix type is identical to `MATSEQDENSE` when constructed with a single process communicator,
180111a5261eSBarry Smith and `MATMPIDENSE` otherwise.
1802209238afSKris Buschelman
18032ef1f0ffSBarry Smith Options Database Key:
180411a5261eSBarry Smith . -mat_type dense - sets the matrix type to `MATDENSE` during a call to `MatSetFromOptions()`
1805209238afSKris Buschelman
1806209238afSKris Buschelman Level: beginner
1807209238afSKris Buschelman
18081cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATSEQDENSE`, `MATMPIDENSE`, `MATDENSECUDA`, `MATDENSEHIP`
18096947451fSStefano Zampini M*/
18106947451fSStefano Zampini
18115d83a8b1SBarry Smith /*@
1812273d9f13SBarry Smith MatMPIDenseSetPreallocation - Sets the array used to store the matrix entries
1813273d9f13SBarry Smith
18148f490ea3SStefano Zampini Collective
1815273d9f13SBarry Smith
1816273d9f13SBarry Smith Input Parameters:
1817fe59aa6dSJacob Faibussowitsch + B - the matrix
18182ef1f0ffSBarry Smith - data - optional location of matrix data. Set to `NULL` for PETSc
1819273d9f13SBarry Smith to control all matrix memory allocation.
1820273d9f13SBarry Smith
18212ef1f0ffSBarry Smith Level: intermediate
18222ef1f0ffSBarry Smith
1823273d9f13SBarry Smith Notes:
18242ef1f0ffSBarry Smith The dense format is fully compatible with standard Fortran
1825273d9f13SBarry Smith storage by columns.
1826273d9f13SBarry Smith
1827273d9f13SBarry Smith The data input variable is intended primarily for Fortran programmers
1828273d9f13SBarry Smith who wish to allocate their own matrix memory space. Most users should
18292ef1f0ffSBarry Smith set `data` to `NULL`.
1830273d9f13SBarry Smith
18311cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATMPIDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()`
1832273d9f13SBarry Smith @*/
MatMPIDenseSetPreallocation(Mat B,PetscScalar * data)1833d71ae5a4SJacob Faibussowitsch PetscErrorCode MatMPIDenseSetPreallocation(Mat B, PetscScalar *data)
1834d71ae5a4SJacob Faibussowitsch {
1835273d9f13SBarry Smith PetscFunctionBegin;
1836d5ea218eSStefano Zampini PetscValidHeaderSpecific(B, MAT_CLASSID, 1);
1837cac4c232SBarry Smith PetscTryMethod(B, "MatMPIDenseSetPreallocation_C", (Mat, PetscScalar *), (B, data));
18383ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1839273d9f13SBarry Smith }
1840273d9f13SBarry Smith
1841d3042a70SBarry Smith /*@
184211a5261eSBarry Smith MatDensePlaceArray - Allows one to replace the array in a `MATDENSE` matrix with an
1843d3042a70SBarry Smith array provided by the user. This is useful to avoid copying an array
1844d3042a70SBarry Smith into a matrix
1845d3042a70SBarry Smith
1846d3042a70SBarry Smith Not Collective
1847d3042a70SBarry Smith
1848d3042a70SBarry Smith Input Parameters:
1849d3042a70SBarry Smith + mat - the matrix
1850d3042a70SBarry Smith - array - the array in column major order
1851d3042a70SBarry Smith
18522ef1f0ffSBarry Smith Level: developer
18532ef1f0ffSBarry Smith
185411a5261eSBarry Smith Note:
18554d5500e8SJunchao Zhang Adding `const` to `array` was an oversight, see notes in `VecPlaceArray()`.
18564d5500e8SJunchao Zhang
185711a5261eSBarry 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
1858d3042a70SBarry Smith freed when the matrix is destroyed.
1859d3042a70SBarry Smith
18601cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDenseResetArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`,
186111a5261eSBarry Smith `MatDenseReplaceArray()`
1862d3042a70SBarry Smith @*/
MatDensePlaceArray(Mat mat,const PetscScalar * array)1863d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDensePlaceArray(Mat mat, const PetscScalar *array)
1864d71ae5a4SJacob Faibussowitsch {
1865d3042a70SBarry Smith PetscFunctionBegin;
1866d5ea218eSStefano Zampini PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
1867cac4c232SBarry Smith PetscUseMethod(mat, "MatDensePlaceArray_C", (Mat, const PetscScalar *), (mat, array));
18689566063dSJacob Faibussowitsch PetscCall(PetscObjectStateIncrease((PetscObject)mat));
186947d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
1870637a0070SStefano Zampini mat->offloadmask = PETSC_OFFLOAD_CPU;
1871637a0070SStefano Zampini #endif
18723ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1873d3042a70SBarry Smith }
1874d3042a70SBarry Smith
1875d3042a70SBarry Smith /*@
187611a5261eSBarry Smith MatDenseResetArray - Resets the matrix array to that it previously had before the call to `MatDensePlaceArray()`
1877d3042a70SBarry Smith
1878d3042a70SBarry Smith Not Collective
1879d3042a70SBarry Smith
18802fe279fdSBarry Smith Input Parameter:
1881d3042a70SBarry Smith . mat - the matrix
1882d3042a70SBarry Smith
18832ef1f0ffSBarry Smith Level: developer
18842ef1f0ffSBarry Smith
188511a5261eSBarry Smith Note:
188611a5261eSBarry Smith You can only call this after a call to `MatDensePlaceArray()`
1887d3042a70SBarry Smith
18881cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDensePlaceArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`
1889d3042a70SBarry Smith @*/
MatDenseResetArray(Mat mat)1890d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseResetArray(Mat mat)
1891d71ae5a4SJacob Faibussowitsch {
1892d3042a70SBarry Smith PetscFunctionBegin;
1893d5ea218eSStefano Zampini PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
1894cac4c232SBarry Smith PetscUseMethod(mat, "MatDenseResetArray_C", (Mat), (mat));
18959566063dSJacob Faibussowitsch PetscCall(PetscObjectStateIncrease((PetscObject)mat));
18963ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1897d3042a70SBarry Smith }
1898d3042a70SBarry Smith
1899d5ea218eSStefano Zampini /*@
1900d5ea218eSStefano Zampini MatDenseReplaceArray - Allows one to replace the array in a dense matrix with an
1901d5ea218eSStefano Zampini array provided by the user. This is useful to avoid copying an array
1902d5ea218eSStefano Zampini into a matrix
1903d5ea218eSStefano Zampini
1904d5ea218eSStefano Zampini Not Collective
1905d5ea218eSStefano Zampini
1906d5ea218eSStefano Zampini Input Parameters:
1907d5ea218eSStefano Zampini + mat - the matrix
1908d5ea218eSStefano Zampini - array - the array in column major order
1909d5ea218eSStefano Zampini
19102ef1f0ffSBarry Smith Level: developer
19112ef1f0ffSBarry Smith
191211a5261eSBarry Smith Note:
19134d5500e8SJunchao Zhang Adding `const` to `array` was an oversight, see notes in `VecPlaceArray()`.
19144d5500e8SJunchao Zhang
191511a5261eSBarry Smith The memory passed in MUST be obtained with `PetscMalloc()` and CANNOT be
1916d5ea218eSStefano Zampini freed by the user. It will be freed when the matrix is destroyed.
1917d5ea218eSStefano Zampini
19181cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MatDensePlaceArray()`, `MatDenseGetArray()`, `VecReplaceArray()`
1919d5ea218eSStefano Zampini @*/
MatDenseReplaceArray(Mat mat,const PetscScalar * array)1920d71ae5a4SJacob Faibussowitsch PetscErrorCode MatDenseReplaceArray(Mat mat, const PetscScalar *array)
1921d71ae5a4SJacob Faibussowitsch {
1922d5ea218eSStefano Zampini PetscFunctionBegin;
1923d5ea218eSStefano Zampini PetscValidHeaderSpecific(mat, MAT_CLASSID, 1);
1924cac4c232SBarry Smith PetscUseMethod(mat, "MatDenseReplaceArray_C", (Mat, const PetscScalar *), (mat, array));
19259566063dSJacob Faibussowitsch PetscCall(PetscObjectStateIncrease((PetscObject)mat));
192647d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
1927d5ea218eSStefano Zampini mat->offloadmask = PETSC_OFFLOAD_CPU;
1928d5ea218eSStefano Zampini #endif
19293ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
1930d5ea218eSStefano Zampini }
1931d5ea218eSStefano Zampini
19325d83a8b1SBarry Smith /*@
193311a5261eSBarry Smith MatCreateDense - Creates a matrix in `MATDENSE` format.
19348965ea79SLois Curfman McInnes
1935d083f849SBarry Smith Collective
1936db81eaa0SLois Curfman McInnes
19378965ea79SLois Curfman McInnes Input Parameters:
1938db81eaa0SLois Curfman McInnes + comm - MPI communicator
19392ef1f0ffSBarry Smith . m - number of local rows (or `PETSC_DECIDE` to have calculated if `M` is given)
19402ef1f0ffSBarry Smith . n - number of local columns (or `PETSC_DECIDE` to have calculated if `N` is given)
19412ef1f0ffSBarry Smith . M - number of global rows (or `PETSC_DECIDE` to have calculated if `m` is given)
19422ef1f0ffSBarry Smith . N - number of global columns (or `PETSC_DECIDE` to have calculated if `n` is given)
19432c5cda21SPierre Jolivet - data - optional location of matrix data. Set data to `NULL` (`PETSC_NULL_SCALAR_ARRAY` for Fortran users) for PETSc
1944dfc5480cSLois Curfman McInnes to control all matrix memory allocation.
19458965ea79SLois Curfman McInnes
19468965ea79SLois Curfman McInnes Output Parameter:
1947477f1c0bSLois Curfman McInnes . A - the matrix
19488965ea79SLois Curfman McInnes
19492ef1f0ffSBarry Smith Level: intermediate
19502ef1f0ffSBarry Smith
1951b259b22eSLois Curfman McInnes Notes:
19522ef1f0ffSBarry Smith The dense format is fully compatible with standard Fortran
195339ddd567SLois Curfman McInnes storage by columns.
195411a5261eSBarry Smith
195511a5261eSBarry Smith Although local portions of the matrix are stored in column-major
1956002bc6daSRichard Tran Mills order, the matrix is partitioned across MPI ranks by row.
19578965ea79SLois Curfman McInnes
195818f449edSLois Curfman McInnes The data input variable is intended primarily for Fortran programmers
195918f449edSLois Curfman McInnes who wish to allocate their own matrix memory space. Most users should
19602c5cda21SPierre Jolivet set `data` to `NULL` (`PETSC_NULL_SCALAR_ARRAY` for Fortran users).
196118f449edSLois Curfman McInnes
19628965ea79SLois Curfman McInnes The user MUST specify either the local or global matrix dimensions
19638965ea79SLois Curfman McInnes (possibly both).
19648965ea79SLois Curfman McInnes
19651cc06b55SBarry Smith .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()`
19668965ea79SLois Curfman McInnes @*/
MatCreateDense(MPI_Comm comm,PetscInt m,PetscInt n,PetscInt M,PetscInt N,PetscScalar data[],Mat * A)1967ce78bad3SBarry Smith PetscErrorCode MatCreateDense(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt M, PetscInt N, PetscScalar data[], Mat *A)
1968d71ae5a4SJacob Faibussowitsch {
19693a40ed3dSBarry Smith PetscFunctionBegin;
19709566063dSJacob Faibussowitsch PetscCall(MatCreate(comm, A));
19719566063dSJacob Faibussowitsch PetscCall(MatSetSizes(*A, m, n, M, N));
19726f08ad05SPierre Jolivet PetscCall(MatSetType(*A, MATDENSE));
19739566063dSJacob Faibussowitsch PetscCall(MatSeqDenseSetPreallocation(*A, data));
19746f08ad05SPierre Jolivet PetscCall(MatMPIDenseSetPreallocation(*A, data));
19753ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
19768965ea79SLois Curfman McInnes }
19778965ea79SLois Curfman McInnes
MatDuplicate_MPIDense(Mat A,MatDuplicateOption cpvalues,Mat * newmat)1978d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatDuplicate_MPIDense(Mat A, MatDuplicateOption cpvalues, Mat *newmat)
1979d71ae5a4SJacob Faibussowitsch {
19808965ea79SLois Curfman McInnes Mat mat;
19813501a2bdSLois Curfman McInnes Mat_MPIDense *a, *oldmat = (Mat_MPIDense *)A->data;
19828965ea79SLois Curfman McInnes
19833a40ed3dSBarry Smith PetscFunctionBegin;
1984f4259b30SLisandro Dalcin *newmat = NULL;
19859566063dSJacob Faibussowitsch PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat));
19869566063dSJacob Faibussowitsch PetscCall(MatSetSizes(mat, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N));
19879566063dSJacob Faibussowitsch PetscCall(MatSetType(mat, ((PetscObject)A)->type_name));
1988834f8fabSBarry Smith a = (Mat_MPIDense *)mat->data;
19895aa7edbeSHong Zhang
1990d5f3da31SBarry Smith mat->factortype = A->factortype;
1991c456f294SBarry Smith mat->assembled = PETSC_TRUE;
1992273d9f13SBarry Smith mat->preallocated = PETSC_TRUE;
19938965ea79SLois Curfman McInnes
1994e0fa3b82SLois Curfman McInnes mat->insertmode = NOT_SET_VALUES;
19953782ba37SSatish Balay a->donotstash = oldmat->donotstash;
1996e04c1aa4SHong Zhang
19979566063dSJacob Faibussowitsch PetscCall(PetscLayoutReference(A->rmap, &mat->rmap));
19989566063dSJacob Faibussowitsch PetscCall(PetscLayoutReference(A->cmap, &mat->cmap));
19998965ea79SLois Curfman McInnes
20009566063dSJacob Faibussowitsch PetscCall(MatDuplicate(oldmat->A, cpvalues, &a->A));
200101b82886SBarry Smith
20028965ea79SLois Curfman McInnes *newmat = mat;
20033ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
20048965ea79SLois Curfman McInnes }
20058965ea79SLois Curfman McInnes
MatLoad_MPIDense(Mat newMat,PetscViewer viewer)200666976f2fSJacob Faibussowitsch static PetscErrorCode MatLoad_MPIDense(Mat newMat, PetscViewer viewer)
2007d71ae5a4SJacob Faibussowitsch {
200887d5ce66SSatish Balay PetscBool isbinary;
200987d5ce66SSatish Balay #if defined(PETSC_HAVE_HDF5)
201087d5ce66SSatish Balay PetscBool ishdf5;
201187d5ce66SSatish Balay #endif
2012eb91f321SVaclav Hapla
2013eb91f321SVaclav Hapla PetscFunctionBegin;
2014eb91f321SVaclav Hapla PetscValidHeaderSpecific(newMat, MAT_CLASSID, 1);
2015eb91f321SVaclav Hapla PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
2016eb91f321SVaclav Hapla /* force binary viewer to load .info file if it has not yet done so */
20179566063dSJacob Faibussowitsch PetscCall(PetscViewerSetUp(viewer));
20189566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
201987d5ce66SSatish Balay #if defined(PETSC_HAVE_HDF5)
20209566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERHDF5, &ishdf5));
202187d5ce66SSatish Balay #endif
2022eb91f321SVaclav Hapla if (isbinary) {
20239566063dSJacob Faibussowitsch PetscCall(MatLoad_Dense_Binary(newMat, viewer));
2024eb91f321SVaclav Hapla #if defined(PETSC_HAVE_HDF5)
202587d5ce66SSatish Balay } else if (ishdf5) {
20269566063dSJacob Faibussowitsch PetscCall(MatLoad_Dense_HDF5(newMat, viewer));
2027eb91f321SVaclav Hapla #endif
202898921bdaSJacob 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);
20293ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2030eb91f321SVaclav Hapla }
2031eb91f321SVaclav Hapla
MatEqual_MPIDense(Mat A,Mat B,PetscBool * flag)2032d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatEqual_MPIDense(Mat A, Mat B, PetscBool *flag)
2033d71ae5a4SJacob Faibussowitsch {
20346e4ee0c6SHong Zhang Mat_MPIDense *matB = (Mat_MPIDense *)B->data, *matA = (Mat_MPIDense *)A->data;
20356e4ee0c6SHong Zhang Mat a, b;
203690ace30eSBarry Smith
20376e4ee0c6SHong Zhang PetscFunctionBegin;
20386e4ee0c6SHong Zhang a = matA->A;
20396e4ee0c6SHong Zhang b = matB->A;
20402cf15c64SPierre Jolivet PetscCall(MatEqual(a, b, flag));
20415440e5dcSBarry Smith PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, flag, 1, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
20423ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
20436e4ee0c6SHong Zhang }
204490ace30eSBarry Smith
MatProductCtxDestroy_MatTransMatMult_MPIDense_MPIDense(PetscCtxRt data)20452a8381b2SBarry Smith static PetscErrorCode MatProductCtxDestroy_MatTransMatMult_MPIDense_MPIDense(PetscCtxRt data)
2046d71ae5a4SJacob Faibussowitsch {
2047cc1eb50dSBarry Smith MatProductCtx_TransMatMultDense *atb = *(MatProductCtx_TransMatMultDense **)data;
2048baa3c1c6SHong Zhang
2049baa3c1c6SHong Zhang PetscFunctionBegin;
20509566063dSJacob Faibussowitsch PetscCall(PetscFree2(atb->sendbuf, atb->recvcounts));
20519566063dSJacob Faibussowitsch PetscCall(MatDestroy(&atb->atb));
20529566063dSJacob Faibussowitsch PetscCall(PetscFree(atb));
20533ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2054baa3c1c6SHong Zhang }
2055baa3c1c6SHong Zhang
MatProductCtxDestroy_MatMatTransMult_MPIDense_MPIDense(PetscCtxRt data)20562a8381b2SBarry Smith static PetscErrorCode MatProductCtxDestroy_MatMatTransMult_MPIDense_MPIDense(PetscCtxRt data)
2057d71ae5a4SJacob Faibussowitsch {
2058cc1eb50dSBarry Smith MatProductCtx_MatTransMultDense *abt = *(MatProductCtx_MatTransMultDense **)data;
2059cc48ffa7SToby Isaac
2060cc48ffa7SToby Isaac PetscFunctionBegin;
20619566063dSJacob Faibussowitsch PetscCall(PetscFree2(abt->buf[0], abt->buf[1]));
20629566063dSJacob Faibussowitsch PetscCall(PetscFree2(abt->recvcounts, abt->recvdispls));
20639566063dSJacob Faibussowitsch PetscCall(PetscFree(abt));
20643ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2065cc48ffa7SToby Isaac }
2066cc48ffa7SToby Isaac
MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat A,Mat B,Mat C)2067d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2068d71ae5a4SJacob Faibussowitsch {
2069baa3c1c6SHong Zhang Mat_MPIDense *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
2070cc1eb50dSBarry Smith MatProductCtx_TransMatMultDense *atb;
2071cb20be35SHong Zhang MPI_Comm comm;
20726718818eSStefano Zampini PetscMPIInt size, *recvcounts;
20736718818eSStefano Zampini PetscScalar *carray, *sendbuf;
2074637a0070SStefano Zampini const PetscScalar *atbarray;
2075a1f9a5e6SStefano Zampini PetscInt i, cN = C->cmap->N, proc, k, j, lda;
2076e68c0b26SHong Zhang const PetscInt *ranges;
2077cb20be35SHong Zhang
2078cb20be35SHong Zhang PetscFunctionBegin;
20796718818eSStefano Zampini MatCheckProduct(C, 3);
20805f80ce2aSJacob Faibussowitsch PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2081cc1eb50dSBarry Smith atb = (MatProductCtx_TransMatMultDense *)C->product->data;
20826718818eSStefano Zampini recvcounts = atb->recvcounts;
20836718818eSStefano Zampini sendbuf = atb->sendbuf;
20846718818eSStefano Zampini
20859566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
20869566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm, &size));
2087e68c0b26SHong Zhang
2088c5ef1628SHong Zhang /* compute atbarray = aseq^T * bseq */
2089fb842aefSJose E. Roman PetscCall(MatTransposeMatMult(a->A, b->A, atb->atb ? MAT_REUSE_MATRIX : MAT_INITIAL_MATRIX, PETSC_DETERMINE, &atb->atb));
2090cb20be35SHong Zhang
20919566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRanges(C, &ranges));
2092cb20be35SHong Zhang
2093d016bddeSToby Isaac if (ranges[1] == C->rmap->N) {
2094d016bddeSToby Isaac /* all of the values are being reduced to rank 0: optimize this case to use MPI_Reduce and GPU aware MPI if available */
2095d016bddeSToby Isaac PetscInt atb_lda, c_lda;
2096d016bddeSToby Isaac Mat atb_local = atb->atb;
2097d016bddeSToby Isaac Mat atb_alloc = NULL;
2098d016bddeSToby Isaac Mat c_local = c->A;
2099d016bddeSToby Isaac Mat c_alloc = NULL;
2100d016bddeSToby Isaac PetscMemType atb_memtype, c_memtype;
2101d016bddeSToby Isaac const PetscScalar *atb_array = NULL;
2102d016bddeSToby Isaac MPI_Datatype vector_type;
2103d016bddeSToby Isaac PetscScalar *c_array = NULL;
2104d016bddeSToby Isaac PetscMPIInt rank;
2105d016bddeSToby Isaac
2106d016bddeSToby Isaac PetscCallMPI(MPI_Comm_rank(comm, &rank));
2107d016bddeSToby Isaac
2108d016bddeSToby Isaac PetscCall(MatDenseGetLDA(atb_local, &atb_lda));
2109d016bddeSToby Isaac if (atb_lda != C->rmap->N) {
2110d016bddeSToby Isaac // copy atb to a matrix that will have lda == the number of rows
2111d016bddeSToby Isaac PetscCall(MatDuplicate(atb_local, MAT_DO_NOT_COPY_VALUES, &atb_alloc));
2112d016bddeSToby Isaac PetscCall(MatCopy(atb_local, atb_alloc, DIFFERENT_NONZERO_PATTERN));
2113d016bddeSToby Isaac atb_local = atb_alloc;
2114d016bddeSToby Isaac }
2115d016bddeSToby Isaac
2116d016bddeSToby Isaac if (rank == 0) {
2117d016bddeSToby Isaac PetscCall(MatDenseGetLDA(c_local, &c_lda));
2118d016bddeSToby Isaac if (c_lda != C->rmap->N) {
2119d016bddeSToby Isaac // copy c to a matrix that will have lda == the number of rows
2120d016bddeSToby Isaac PetscCall(MatDuplicate(c_local, MAT_DO_NOT_COPY_VALUES, &c_alloc));
2121d016bddeSToby Isaac c_local = c_alloc;
2122d016bddeSToby Isaac }
2123d016bddeSToby Isaac PetscCall(MatZeroEntries(c_local));
2124d016bddeSToby Isaac }
2125d016bddeSToby Isaac /* atb_local and c_local have nrows = lda = A->cmap->N and ncols =
2126d016bddeSToby Isaac * B->cmap->N: use the a->Mvctx to use the best reduction method */
2127d016bddeSToby Isaac if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
2128d016bddeSToby Isaac vector_type = MPIU_SCALAR;
2129d016bddeSToby Isaac if (B->cmap->N > 1) {
2130b682a77bSToby Isaac PetscMPIInt mpi_N;
2131b682a77bSToby Isaac
2132b682a77bSToby Isaac PetscCall(PetscMPIIntCast(B->cmap->N, &mpi_N));
2133b682a77bSToby Isaac PetscCallMPI(MPI_Type_contiguous(mpi_N, MPIU_SCALAR, &vector_type));
2134d016bddeSToby Isaac PetscCallMPI(MPI_Type_commit(&vector_type));
2135d016bddeSToby Isaac }
2136d016bddeSToby Isaac PetscCall(MatDenseGetArrayReadAndMemType(atb_local, &atb_array, &atb_memtype));
2137d016bddeSToby Isaac PetscCall(MatDenseGetArrayWriteAndMemType(c_local, &c_array, &c_memtype));
2138d016bddeSToby Isaac PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, vector_type, atb_memtype, atb_array, c_memtype, c_array, MPIU_SUM));
2139d016bddeSToby Isaac PetscCall(PetscSFReduceEnd(a->Mvctx, vector_type, atb_array, c_array, MPIU_SUM));
2140d016bddeSToby Isaac PetscCall(MatDenseRestoreArrayWriteAndMemType(c_local, &c_array));
2141d016bddeSToby Isaac PetscCall(MatDenseRestoreArrayReadAndMemType(atb_local, &atb_array));
2142d016bddeSToby Isaac if (rank == 0 && c_local != c->A) PetscCall(MatCopy(c_local, c->A, DIFFERENT_NONZERO_PATTERN));
2143d016bddeSToby Isaac if (B->cmap->N > 1) PetscCallMPI(MPI_Type_free(&vector_type));
2144d016bddeSToby Isaac PetscCall(MatDestroy(&atb_alloc));
2145d016bddeSToby Isaac PetscCall(MatDestroy(&c_alloc));
2146d016bddeSToby Isaac PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2147d016bddeSToby Isaac PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2148d016bddeSToby Isaac PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2149d016bddeSToby Isaac PetscFunctionReturn(PETSC_SUCCESS);
2150d016bddeSToby Isaac }
2151d016bddeSToby Isaac
2152660d5466SHong Zhang /* arrange atbarray into sendbuf */
21539566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(atb->atb, &atbarray));
21549566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(atb->atb, &lda));
2155637a0070SStefano Zampini for (proc = 0, k = 0; proc < size; proc++) {
2156baa3c1c6SHong Zhang for (j = 0; j < cN; j++) {
2157a1f9a5e6SStefano Zampini for (i = ranges[proc]; i < ranges[proc + 1]; i++) sendbuf[k++] = atbarray[i + j * lda];
2158cb20be35SHong Zhang }
2159cb20be35SHong Zhang }
21609566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(atb->atb, &atbarray));
2161637a0070SStefano Zampini
2162c5ef1628SHong Zhang /* sum all atbarray to local values of C */
21639566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayWrite(c->A, &carray));
21649566063dSJacob Faibussowitsch PetscCallMPI(MPI_Reduce_scatter(sendbuf, carray, recvcounts, MPIU_SCALAR, MPIU_SUM, comm));
21659566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayWrite(c->A, &carray));
216667af85e8SPierre Jolivet PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
21679566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
21689566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
21693ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2170cb20be35SHong Zhang }
2171cb20be35SHong Zhang
MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat A,Mat B,PetscReal fill,Mat C)2172d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2173d71ae5a4SJacob Faibussowitsch {
2174cb20be35SHong Zhang MPI_Comm comm;
2175baa3c1c6SHong Zhang PetscMPIInt size;
2176660d5466SHong Zhang PetscInt cm = A->cmap->n, cM, cN = B->cmap->N;
2177cc1eb50dSBarry Smith MatProductCtx_TransMatMultDense *atb;
217847d993e7Ssuyashtn PetscBool cisdense = PETSC_FALSE;
21790acaeb71SStefano Zampini const PetscInt *ranges;
2180cb20be35SHong Zhang
2181cb20be35SHong Zhang PetscFunctionBegin;
21828f83a4d9SJacob Faibussowitsch MatCheckProduct(C, 4);
21835f80ce2aSJacob Faibussowitsch PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
21849566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
2185966bd95aSPierre Jolivet PetscCheck(A->rmap->rstart == B->rmap->rstart && A->rmap->rend == B->rmap->rend, comm, PETSC_ERR_ARG_SIZ, "Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != B (%" PetscInt_FMT ",%" PetscInt_FMT ")", A->rmap->rstart,
2186966bd95aSPierre Jolivet A->rmap->rend, B->rmap->rstart, B->rmap->rend);
2187cb20be35SHong Zhang
21884222ddf1SHong Zhang /* create matrix product C */
21899566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C, cm, B->cmap->n, A->cmap->N, B->cmap->N));
219047d993e7Ssuyashtn #if defined(PETSC_HAVE_CUDA)
21919566063dSJacob Faibussowitsch PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSECUDA, ""));
219247d993e7Ssuyashtn #endif
219347d993e7Ssuyashtn #if defined(PETSC_HAVE_HIP)
219447d993e7Ssuyashtn PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSEHIP, ""));
219547d993e7Ssuyashtn #endif
219648a46eb9SPierre Jolivet if (!cisdense) PetscCall(MatSetType(C, ((PetscObject)A)->type_name));
21979566063dSJacob Faibussowitsch PetscCall(MatSetUp(C));
2198baa3c1c6SHong Zhang
21994222ddf1SHong Zhang /* create data structure for reuse C */
22009566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm, &size));
22019566063dSJacob Faibussowitsch PetscCall(PetscNew(&atb));
22024222ddf1SHong Zhang cM = C->rmap->N;
22039566063dSJacob Faibussowitsch PetscCall(PetscMalloc2(cM * cN, &atb->sendbuf, size, &atb->recvcounts));
22049566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRanges(C, &ranges));
22056497c311SBarry Smith for (PetscMPIInt i = 0; i < size; i++) PetscCall(PetscMPIIntCast((ranges[i + 1] - ranges[i]) * cN, &atb->recvcounts[i]));
22066718818eSStefano Zampini C->product->data = atb;
2207cc1eb50dSBarry Smith C->product->destroy = MatProductCtxDestroy_MatTransMatMult_MPIDense_MPIDense;
22083ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2209cb20be35SHong Zhang }
2210cb20be35SHong Zhang
MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat A,Mat B,PetscReal fill,Mat C)2211d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2212d71ae5a4SJacob Faibussowitsch {
2213cc48ffa7SToby Isaac MPI_Comm comm;
2214cc48ffa7SToby Isaac PetscMPIInt i, size;
2215cc48ffa7SToby Isaac PetscInt maxRows, bufsiz;
2216cc48ffa7SToby Isaac PetscMPIInt tag;
22174222ddf1SHong Zhang PetscInt alg;
2218cc1eb50dSBarry Smith MatProductCtx_MatTransMultDense *abt;
22194222ddf1SHong Zhang Mat_Product *product = C->product;
22204222ddf1SHong Zhang PetscBool flg;
2221cc48ffa7SToby Isaac
2222cc48ffa7SToby Isaac PetscFunctionBegin;
22236718818eSStefano Zampini MatCheckProduct(C, 4);
22245f80ce2aSJacob Faibussowitsch PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
22254222ddf1SHong Zhang /* check local size of A and B */
22265f80ce2aSJacob 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);
2227cc48ffa7SToby Isaac
22289566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg, "allgatherv", &flg));
2229637a0070SStefano Zampini alg = flg ? 0 : 1;
2230cc48ffa7SToby Isaac
22314222ddf1SHong Zhang /* setup matrix product C */
22329566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C, A->rmap->n, B->rmap->n, A->rmap->N, B->rmap->N));
22339566063dSJacob Faibussowitsch PetscCall(MatSetType(C, MATMPIDENSE));
22349566063dSJacob Faibussowitsch PetscCall(MatSetUp(C));
22359566063dSJacob Faibussowitsch PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag));
22364222ddf1SHong Zhang
22374222ddf1SHong Zhang /* create data structure for reuse C */
22389566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
22399566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm, &size));
22409566063dSJacob Faibussowitsch PetscCall(PetscNew(&abt));
2241cc48ffa7SToby Isaac abt->tag = tag;
2242faa55883SToby Isaac abt->alg = alg;
2243faa55883SToby Isaac switch (alg) {
22444222ddf1SHong Zhang case 1: /* alg: "cyclic" */
224557508eceSPierre Jolivet for (maxRows = 0, i = 0; i < size; i++) maxRows = PetscMax(maxRows, B->rmap->range[i + 1] - B->rmap->range[i]);
2246cc48ffa7SToby Isaac bufsiz = A->cmap->N * maxRows;
2247f4f49eeaSPierre Jolivet PetscCall(PetscMalloc2(bufsiz, &abt->buf[0], bufsiz, &abt->buf[1]));
2248faa55883SToby Isaac break;
22494222ddf1SHong Zhang default: /* alg: "allgatherv" */
2250f4f49eeaSPierre Jolivet PetscCall(PetscMalloc2(B->rmap->n * B->cmap->N, &abt->buf[0], B->rmap->N * B->cmap->N, &abt->buf[1]));
2251f4f49eeaSPierre Jolivet PetscCall(PetscMalloc2(size, &abt->recvcounts, size + 1, &abt->recvdispls));
22526497c311SBarry Smith for (i = 0; i <= size; i++) PetscCall(PetscMPIIntCast(B->rmap->range[i] * A->cmap->N, &abt->recvdispls[i]));
22536497c311SBarry Smith for (i = 0; i < size; i++) PetscCall(PetscMPIIntCast(abt->recvdispls[i + 1] - abt->recvdispls[i], &abt->recvcounts[i]));
2254faa55883SToby Isaac break;
2255faa55883SToby Isaac }
2256cc48ffa7SToby Isaac
22576718818eSStefano Zampini C->product->data = abt;
2258cc1eb50dSBarry Smith C->product->destroy = MatProductCtxDestroy_MatMatTransMult_MPIDense_MPIDense;
22594222ddf1SHong Zhang C->ops->mattransposemultnumeric = MatMatTransposeMultNumeric_MPIDense_MPIDense;
22603ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2261cc48ffa7SToby Isaac }
2262cc48ffa7SToby Isaac
MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(Mat A,Mat B,Mat C)2263d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(Mat A, Mat B, Mat C)
2264d71ae5a4SJacob Faibussowitsch {
2265cc48ffa7SToby Isaac Mat_MPIDense *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
2266cc1eb50dSBarry Smith MatProductCtx_MatTransMultDense *abt;
2267cc48ffa7SToby Isaac MPI_Comm comm;
22686497c311SBarry Smith PetscMPIInt rank, size, sendto, recvfrom, recvisfrom;
2269f4259b30SLisandro Dalcin PetscScalar *sendbuf, *recvbuf = NULL, *cv;
22706497c311SBarry Smith PetscInt i, cK = A->cmap->N, sendsiz, recvsiz, k, j, bn;
2271cc48ffa7SToby Isaac PetscScalar _DOne = 1.0, _DZero = 0.0;
2272637a0070SStefano Zampini const PetscScalar *av, *bv;
227323fff9afSBarry Smith PetscBLASInt cm, cn, ck, alda, blda = 0, clda;
2274cc48ffa7SToby Isaac MPI_Request reqs[2];
2275cc48ffa7SToby Isaac const PetscInt *ranges;
2276cc48ffa7SToby Isaac
2277cc48ffa7SToby Isaac PetscFunctionBegin;
22786718818eSStefano Zampini MatCheckProduct(C, 3);
22795f80ce2aSJacob Faibussowitsch PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2280cc1eb50dSBarry Smith abt = (MatProductCtx_MatTransMultDense *)C->product->data;
22819566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
22829566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm, &rank));
22839566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm, &size));
22849566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(a->A, &av));
22859566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(b->A, &bv));
22869566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayWrite(c->A, &cv));
22879566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &i));
22889566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(i, &alda));
22899566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(b->A, &i));
22909566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(i, &blda));
22919566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(c->A, &i));
22929566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(i, &clda));
22939566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRanges(B, &ranges));
2294cc48ffa7SToby Isaac bn = B->rmap->n;
2295637a0070SStefano Zampini if (blda == bn) {
2296637a0070SStefano Zampini sendbuf = (PetscScalar *)bv;
2297cc48ffa7SToby Isaac } else {
2298cc48ffa7SToby Isaac sendbuf = abt->buf[0];
2299cc48ffa7SToby Isaac for (k = 0, i = 0; i < cK; i++) {
2300ad540459SPierre Jolivet for (j = 0; j < bn; j++, k++) sendbuf[k] = bv[i * blda + j];
2301cc48ffa7SToby Isaac }
2302cc48ffa7SToby Isaac }
2303cc48ffa7SToby Isaac if (size > 1) {
2304cc48ffa7SToby Isaac sendto = (rank + size - 1) % size;
2305cc48ffa7SToby Isaac recvfrom = (rank + size + 1) % size;
2306cc48ffa7SToby Isaac } else {
2307cc48ffa7SToby Isaac sendto = recvfrom = 0;
2308cc48ffa7SToby Isaac }
23099566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(cK, &ck));
23109566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm));
2311cc48ffa7SToby Isaac recvisfrom = rank;
2312cc48ffa7SToby Isaac for (i = 0; i < size; i++) {
2313cc48ffa7SToby Isaac /* we have finished receiving in sending, bufs can be read/modified */
23146497c311SBarry Smith PetscMPIInt nextrecvisfrom = (recvisfrom + 1) % size; /* which process the next recvbuf will originate on */
2315cc48ffa7SToby Isaac PetscInt nextbn = ranges[nextrecvisfrom + 1] - ranges[nextrecvisfrom];
2316cc48ffa7SToby Isaac
2317cc48ffa7SToby Isaac if (nextrecvisfrom != rank) {
2318cc48ffa7SToby Isaac /* start the cyclic sends from sendbuf, to recvbuf (which will switch to sendbuf) */
2319cc48ffa7SToby Isaac sendsiz = cK * bn;
2320cc48ffa7SToby Isaac recvsiz = cK * nextbn;
2321cc48ffa7SToby Isaac recvbuf = (i & 1) ? abt->buf[0] : abt->buf[1];
23226497c311SBarry Smith PetscCallMPI(MPIU_Isend(sendbuf, sendsiz, MPIU_SCALAR, sendto, abt->tag, comm, &reqs[0]));
23236497c311SBarry Smith PetscCallMPI(MPIU_Irecv(recvbuf, recvsiz, MPIU_SCALAR, recvfrom, abt->tag, comm, &reqs[1]));
2324cc48ffa7SToby Isaac }
2325cc48ffa7SToby Isaac
2326cc48ffa7SToby Isaac /* local aseq * sendbuf^T */
23279566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(ranges[recvisfrom + 1] - ranges[recvisfrom], &cn));
2328792fecdfSBarry Smith if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "T", &cm, &cn, &ck, &_DOne, av, &alda, sendbuf, &cn, &_DZero, cv + clda * ranges[recvisfrom], &clda));
2329cc48ffa7SToby Isaac
2330cc48ffa7SToby Isaac if (nextrecvisfrom != rank) {
2331cc48ffa7SToby Isaac /* wait for the sends and receives to complete, swap sendbuf and recvbuf */
23329566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitall(2, reqs, MPI_STATUSES_IGNORE));
2333cc48ffa7SToby Isaac }
2334cc48ffa7SToby Isaac bn = nextbn;
2335cc48ffa7SToby Isaac recvisfrom = nextrecvisfrom;
2336cc48ffa7SToby Isaac sendbuf = recvbuf;
2337cc48ffa7SToby Isaac }
23389566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(a->A, &av));
23399566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
23409566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayWrite(c->A, &cv));
234167af85e8SPierre Jolivet PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
23429566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
23439566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
23443ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2345cc48ffa7SToby Isaac }
2346cc48ffa7SToby Isaac
MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(Mat A,Mat B,Mat C)2347d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(Mat A, Mat B, Mat C)
2348d71ae5a4SJacob Faibussowitsch {
2349faa55883SToby Isaac Mat_MPIDense *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
2350cc1eb50dSBarry Smith MatProductCtx_MatTransMultDense *abt;
2351faa55883SToby Isaac MPI_Comm comm;
23526497c311SBarry Smith PetscMPIInt size, ibn;
2353637a0070SStefano Zampini PetscScalar *cv, *sendbuf, *recvbuf;
2354637a0070SStefano Zampini const PetscScalar *av, *bv;
2355637a0070SStefano Zampini PetscInt blda, i, cK = A->cmap->N, k, j, bn;
2356faa55883SToby Isaac PetscScalar _DOne = 1.0, _DZero = 0.0;
2357637a0070SStefano Zampini PetscBLASInt cm, cn, ck, alda, clda;
2358faa55883SToby Isaac
2359faa55883SToby Isaac PetscFunctionBegin;
23606718818eSStefano Zampini MatCheckProduct(C, 3);
23615f80ce2aSJacob Faibussowitsch PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2362cc1eb50dSBarry Smith abt = (MatProductCtx_MatTransMultDense *)C->product->data;
23639566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
23649566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm, &size));
23659566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(a->A, &av));
23669566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(b->A, &bv));
23679566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayWrite(c->A, &cv));
23689566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(a->A, &i));
23699566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(i, &alda));
23709566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(b->A, &blda));
23719566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(c->A, &i));
23729566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(i, &clda));
2373faa55883SToby Isaac /* copy transpose of B into buf[0] */
2374faa55883SToby Isaac bn = B->rmap->n;
2375faa55883SToby Isaac sendbuf = abt->buf[0];
2376faa55883SToby Isaac recvbuf = abt->buf[1];
2377faa55883SToby Isaac for (k = 0, j = 0; j < bn; j++) {
2378ad540459SPierre Jolivet for (i = 0; i < cK; i++, k++) sendbuf[k] = bv[i * blda + j];
2379faa55883SToby Isaac }
23809566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
23816497c311SBarry Smith PetscCall(PetscMPIIntCast(bn * cK, &ibn));
23826497c311SBarry Smith PetscCallMPI(MPI_Allgatherv(sendbuf, ibn, MPIU_SCALAR, recvbuf, abt->recvcounts, abt->recvdispls, MPIU_SCALAR, comm));
23839566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(cK, &ck));
23849566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm));
23859566063dSJacob Faibussowitsch PetscCall(PetscBLASIntCast(c->A->cmap->n, &cn));
2386792fecdfSBarry Smith if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "N", &cm, &cn, &ck, &_DOne, av, &alda, recvbuf, &ck, &_DZero, cv, &clda));
23879566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(a->A, &av));
23889566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
23899566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayWrite(c->A, &cv));
239067af85e8SPierre Jolivet PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
23919566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
23929566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
23933ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2394faa55883SToby Isaac }
2395faa55883SToby Isaac
MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat A,Mat B,Mat C)2396d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2397d71ae5a4SJacob Faibussowitsch {
2398cc1eb50dSBarry Smith MatProductCtx_MatTransMultDense *abt;
2399faa55883SToby Isaac
2400faa55883SToby Isaac PetscFunctionBegin;
24016718818eSStefano Zampini MatCheckProduct(C, 3);
24025f80ce2aSJacob Faibussowitsch PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2403cc1eb50dSBarry Smith abt = (MatProductCtx_MatTransMultDense *)C->product->data;
2404faa55883SToby Isaac switch (abt->alg) {
2405d71ae5a4SJacob Faibussowitsch case 1:
2406d71ae5a4SJacob Faibussowitsch PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(A, B, C));
2407d71ae5a4SJacob Faibussowitsch break;
2408d71ae5a4SJacob Faibussowitsch default:
2409d71ae5a4SJacob Faibussowitsch PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(A, B, C));
2410d71ae5a4SJacob Faibussowitsch break;
2411faa55883SToby Isaac }
24123ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2413faa55883SToby Isaac }
2414faa55883SToby Isaac
MatProductCtxDestroy_MatMatMult_MPIDense_MPIDense(PetscCtxRt data)24152a8381b2SBarry Smith static PetscErrorCode MatProductCtxDestroy_MatMatMult_MPIDense_MPIDense(PetscCtxRt data)
2416d71ae5a4SJacob Faibussowitsch {
2417cc1eb50dSBarry Smith MatProductCtx_MatMultDense *ab = *(MatProductCtx_MatMultDense **)data;
2418320f2790SHong Zhang
2419320f2790SHong Zhang PetscFunctionBegin;
24209566063dSJacob Faibussowitsch PetscCall(MatDestroy(&ab->Ce));
24219566063dSJacob Faibussowitsch PetscCall(MatDestroy(&ab->Ae));
24229566063dSJacob Faibussowitsch PetscCall(MatDestroy(&ab->Be));
24239566063dSJacob Faibussowitsch PetscCall(PetscFree(ab));
24243ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2425320f2790SHong Zhang }
2426320f2790SHong Zhang
MatMatMultNumeric_MPIDense_MPIDense(Mat A,Mat B,Mat C)24275d8c7819SPierre Jolivet static PetscErrorCode MatMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2428d71ae5a4SJacob Faibussowitsch {
2429cc1eb50dSBarry Smith MatProductCtx_MatMultDense *ab;
24305d8c7819SPierre Jolivet Mat_MPIDense *mdn = (Mat_MPIDense *)A->data;
2431d016bddeSToby Isaac Mat_MPIDense *b = (Mat_MPIDense *)B->data;
2432320f2790SHong Zhang
2433320f2790SHong Zhang PetscFunctionBegin;
24346718818eSStefano Zampini MatCheckProduct(C, 3);
24355f80ce2aSJacob Faibussowitsch PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Missing product data");
2436cc1eb50dSBarry Smith ab = (MatProductCtx_MatMultDense *)C->product->data;
24375d8c7819SPierre Jolivet if (ab->Ae && ab->Ce) {
24385d8c7819SPierre Jolivet #if PetscDefined(HAVE_ELEMENTAL)
24399566063dSJacob Faibussowitsch PetscCall(MatConvert_MPIDense_Elemental(A, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Ae));
24409566063dSJacob Faibussowitsch PetscCall(MatConvert_MPIDense_Elemental(B, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Be));
24419566063dSJacob Faibussowitsch PetscCall(MatMatMultNumeric_Elemental(ab->Ae, ab->Be, ab->Ce));
24429566063dSJacob Faibussowitsch PetscCall(MatConvert(ab->Ce, MATMPIDENSE, MAT_REUSE_MATRIX, &C));
24435d8c7819SPierre Jolivet #else
24445d8c7819SPierre Jolivet SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined");
24455d8c7819SPierre Jolivet #endif
24465d8c7819SPierre Jolivet } else {
2447d016bddeSToby Isaac MPI_Comm comm;
24485d8c7819SPierre Jolivet const PetscScalar *read;
24495d8c7819SPierre Jolivet PetscScalar *write;
24505d8c7819SPierre Jolivet PetscInt lda;
2451d016bddeSToby Isaac const PetscInt *ranges;
2452d016bddeSToby Isaac PetscMPIInt size;
24535d8c7819SPierre Jolivet
2454d016bddeSToby Isaac if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A)); /* cannot be done during the symbolic phase because of possible calls to MatProductReplaceMats() */
2455d016bddeSToby Isaac comm = PetscObjectComm((PetscObject)B);
2456d016bddeSToby Isaac PetscCallMPI(MPI_Comm_size(comm, &size));
2457d016bddeSToby Isaac PetscCall(PetscLayoutGetRanges(B->rmap, &ranges));
2458d016bddeSToby Isaac if (ranges[1] == ranges[size]) {
2459d016bddeSToby Isaac // optimize for the case where the B matrix is broadcast from rank 0
2460d016bddeSToby Isaac PetscInt b_lda, be_lda;
2461d016bddeSToby Isaac Mat b_local = b->A;
2462d016bddeSToby Isaac Mat b_alloc = NULL;
2463d016bddeSToby Isaac Mat be_local = ab->Be;
2464d016bddeSToby Isaac Mat be_alloc = NULL;
2465d016bddeSToby Isaac PetscMemType b_memtype, be_memtype;
2466d016bddeSToby Isaac const PetscScalar *b_array = NULL;
2467d016bddeSToby Isaac MPI_Datatype vector_type;
2468d016bddeSToby Isaac PetscScalar *be_array = NULL;
2469d016bddeSToby Isaac PetscMPIInt rank;
2470d016bddeSToby Isaac
2471d016bddeSToby Isaac PetscCallMPI(MPI_Comm_rank(comm, &rank));
2472d016bddeSToby Isaac PetscCall(MatDenseGetLDA(be_local, &be_lda));
2473d016bddeSToby Isaac if (be_lda != B->rmap->N) {
2474d016bddeSToby Isaac PetscCall(MatDuplicate(be_local, MAT_DO_NOT_COPY_VALUES, &be_alloc));
2475d016bddeSToby Isaac be_local = be_alloc;
2476d016bddeSToby Isaac }
2477d016bddeSToby Isaac
2478d016bddeSToby Isaac if (rank == 0) {
2479d016bddeSToby Isaac PetscCall(MatDenseGetLDA(b_local, &b_lda));
2480d016bddeSToby Isaac if (b_lda != B->rmap->N) {
2481d016bddeSToby Isaac PetscCall(MatDuplicate(b_local, MAT_DO_NOT_COPY_VALUES, &b_alloc));
2482d016bddeSToby Isaac PetscCall(MatCopy(b_local, b_alloc, DIFFERENT_NONZERO_PATTERN));
2483d016bddeSToby Isaac b_local = b_alloc;
2484d016bddeSToby Isaac }
2485d016bddeSToby Isaac }
2486b682a77bSToby Isaac vector_type = MPIU_SCALAR;
2487b682a77bSToby Isaac if (B->cmap->N > 1) {
2488b682a77bSToby Isaac PetscMPIInt mpi_N;
2489b682a77bSToby Isaac
2490b682a77bSToby Isaac PetscCall(PetscMPIIntCast(B->cmap->N, &mpi_N));
2491b682a77bSToby Isaac PetscCallMPI(MPI_Type_contiguous(mpi_N, MPIU_SCALAR, &vector_type));
2492d016bddeSToby Isaac PetscCallMPI(MPI_Type_commit(&vector_type));
2493b682a77bSToby Isaac }
2494d016bddeSToby Isaac PetscCall(MatDenseGetArrayReadAndMemType(b_local, &b_array, &b_memtype));
2495d016bddeSToby Isaac PetscCall(MatDenseGetArrayWriteAndMemType(be_local, &be_array, &be_memtype));
2496d016bddeSToby Isaac PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, vector_type, b_memtype, b_array, be_memtype, be_array, MPI_REPLACE));
2497d016bddeSToby Isaac PetscCall(PetscSFBcastEnd(mdn->Mvctx, vector_type, b_array, be_array, MPI_REPLACE));
2498d016bddeSToby Isaac PetscCall(MatDenseRestoreArrayWriteAndMemType(be_local, &be_array));
2499d016bddeSToby Isaac PetscCall(MatDenseRestoreArrayReadAndMemType(b_local, &b_array));
2500d016bddeSToby Isaac if (be_local != ab->Be) PetscCall(MatCopy(be_local, ab->Be, DIFFERENT_NONZERO_PATTERN));
2501b682a77bSToby Isaac if (B->cmap->N > 1) PetscCallMPI(MPI_Type_free(&vector_type));
2502d016bddeSToby Isaac PetscCall(MatDestroy(&be_alloc));
2503d016bddeSToby Isaac PetscCall(MatDestroy(&b_alloc));
2504d016bddeSToby Isaac } else {
25055d8c7819SPierre Jolivet PetscCall(MatDenseGetLDA(B, &lda));
25065d8c7819SPierre Jolivet PetscCall(MatDenseGetArrayRead(B, &read));
25075d8c7819SPierre Jolivet PetscCall(MatDenseGetArrayWrite(ab->Be, &write));
25085d8c7819SPierre Jolivet for (PetscInt i = 0; i < C->cmap->N; ++i) {
25095d8c7819SPierre Jolivet PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE));
25105d8c7819SPierre Jolivet PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE));
25115d8c7819SPierre Jolivet }
25125d8c7819SPierre Jolivet PetscCall(MatDenseRestoreArrayWrite(ab->Be, &write));
25135d8c7819SPierre Jolivet PetscCall(MatDenseRestoreArrayRead(B, &read));
2514d016bddeSToby Isaac }
2515f4f49eeaSPierre Jolivet PetscCall(MatMatMultNumeric_SeqDense_SeqDense(((Mat_MPIDense *)A->data)->A, ab->Be, ((Mat_MPIDense *)C->data)->A));
25165d8c7819SPierre Jolivet }
25173ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2518320f2790SHong Zhang }
2519320f2790SHong Zhang
MatMatMultSymbolic_MPIDense_MPIDense(Mat A,Mat B,PetscReal fill,Mat C)2520d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2521d71ae5a4SJacob Faibussowitsch {
25225d8c7819SPierre Jolivet Mat_Product *product = C->product;
25235d8c7819SPierre Jolivet PetscInt alg;
2524cc1eb50dSBarry Smith MatProductCtx_MatMultDense *ab;
25255d8c7819SPierre Jolivet PetscBool flg;
2526320f2790SHong Zhang
2527320f2790SHong Zhang PetscFunctionBegin;
25286718818eSStefano Zampini MatCheckProduct(C, 4);
25295f80ce2aSJacob Faibussowitsch PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
25304222ddf1SHong Zhang /* check local size of A and B */
25315d8c7819SPierre 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 ")",
25325d8c7819SPierre Jolivet A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
2533320f2790SHong Zhang
25345d8c7819SPierre Jolivet PetscCall(PetscStrcmp(product->alg, "petsc", &flg));
25355d8c7819SPierre Jolivet alg = flg ? 0 : 1;
25364222ddf1SHong Zhang
25374222ddf1SHong Zhang /* setup C */
25385d8c7819SPierre Jolivet PetscCall(MatSetSizes(C, A->rmap->n, B->cmap->n, A->rmap->N, B->cmap->N));
25395d8c7819SPierre Jolivet PetscCall(MatSetType(C, MATMPIDENSE));
25409566063dSJacob Faibussowitsch PetscCall(MatSetUp(C));
2541320f2790SHong Zhang
2542320f2790SHong Zhang /* create data structure for reuse Cdense */
25439566063dSJacob Faibussowitsch PetscCall(PetscNew(&ab));
25445d8c7819SPierre Jolivet
25455d8c7819SPierre Jolivet switch (alg) {
25465d8c7819SPierre Jolivet case 1: /* alg: "elemental" */
25475d8c7819SPierre Jolivet #if PetscDefined(HAVE_ELEMENTAL)
25485d8c7819SPierre Jolivet /* create elemental matrices Ae and Be */
25495d8c7819SPierre Jolivet PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &ab->Ae));
25505d8c7819SPierre Jolivet PetscCall(MatSetSizes(ab->Ae, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N));
25515d8c7819SPierre Jolivet PetscCall(MatSetType(ab->Ae, MATELEMENTAL));
25525d8c7819SPierre Jolivet PetscCall(MatSetUp(ab->Ae));
25535d8c7819SPierre Jolivet PetscCall(MatSetOption(ab->Ae, MAT_ROW_ORIENTED, PETSC_FALSE));
25545d8c7819SPierre Jolivet
25555d8c7819SPierre Jolivet PetscCall(MatCreate(PetscObjectComm((PetscObject)B), &ab->Be));
25565d8c7819SPierre Jolivet PetscCall(MatSetSizes(ab->Be, PETSC_DECIDE, PETSC_DECIDE, B->rmap->N, B->cmap->N));
25575d8c7819SPierre Jolivet PetscCall(MatSetType(ab->Be, MATELEMENTAL));
25585d8c7819SPierre Jolivet PetscCall(MatSetUp(ab->Be));
25595d8c7819SPierre Jolivet PetscCall(MatSetOption(ab->Be, MAT_ROW_ORIENTED, PETSC_FALSE));
25605d8c7819SPierre Jolivet
25615d8c7819SPierre Jolivet /* compute symbolic Ce = Ae*Be */
25625d8c7819SPierre Jolivet PetscCall(MatCreate(PetscObjectComm((PetscObject)C), &ab->Ce));
25635d8c7819SPierre Jolivet PetscCall(MatMatMultSymbolic_Elemental(ab->Ae, ab->Be, fill, ab->Ce));
25645d8c7819SPierre Jolivet #else
25655d8c7819SPierre Jolivet SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined");
25665d8c7819SPierre Jolivet #endif
25675d8c7819SPierre Jolivet break;
25685d8c7819SPierre Jolivet default: /* alg: "petsc" */
25695d8c7819SPierre Jolivet ab->Ae = NULL;
25705d8c7819SPierre Jolivet PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, A->cmap->N, B->cmap->N, NULL, &ab->Be));
25715d8c7819SPierre Jolivet ab->Ce = NULL;
25725d8c7819SPierre Jolivet break;
25735d8c7819SPierre Jolivet }
25746718818eSStefano Zampini
25756718818eSStefano Zampini C->product->data = ab;
2576cc1eb50dSBarry Smith C->product->destroy = MatProductCtxDestroy_MatMatMult_MPIDense_MPIDense;
25774222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIDense_MPIDense;
25783ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2579320f2790SHong Zhang }
25802ef1f0ffSBarry Smith
MatProductSetFromOptions_MPIDense_AB(Mat C)2581d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_AB(Mat C)
2582d71ae5a4SJacob Faibussowitsch {
25835d8c7819SPierre Jolivet Mat_Product *product = C->product;
25845d8c7819SPierre Jolivet const char *algTypes[2] = {"petsc", "elemental"};
25855d8c7819SPierre Jolivet PetscInt alg, nalg = PetscDefined(HAVE_ELEMENTAL) ? 2 : 1;
25865d8c7819SPierre Jolivet PetscBool flg = PETSC_FALSE;
25875d8c7819SPierre Jolivet
25884d86920dSPierre Jolivet PetscFunctionBegin;
25895d8c7819SPierre Jolivet /* Set default algorithm */
2590f0b74427SPierre Jolivet alg = 0; /* default is PETSc */
25915d8c7819SPierre Jolivet PetscCall(PetscStrcmp(product->alg, "default", &flg));
2592835f2295SStefano Zampini if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
25935d8c7819SPierre Jolivet
25945d8c7819SPierre Jolivet /* Get runtime option */
25955d8c7819SPierre Jolivet PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_AB", "Mat");
25965d8c7819SPierre Jolivet PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_AB", algTypes, nalg, algTypes[alg], &alg, &flg));
25975d8c7819SPierre Jolivet PetscOptionsEnd();
2598835f2295SStefano Zampini if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
25995d8c7819SPierre Jolivet
26004222ddf1SHong Zhang C->ops->matmultsymbolic = MatMatMultSymbolic_MPIDense_MPIDense;
26014222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AB;
26023ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
2603320f2790SHong Zhang }
260486aefd0dSHong Zhang
MatProductSetFromOptions_MPIDense_AtB(Mat C)2605d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_AtB(Mat C)
2606d71ae5a4SJacob Faibussowitsch {
26074222ddf1SHong Zhang Mat_Product *product = C->product;
26084222ddf1SHong Zhang Mat A = product->A, B = product->B;
26094222ddf1SHong Zhang
26104222ddf1SHong Zhang PetscFunctionBegin;
261107e90273SPierre 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 ")",
261207e90273SPierre Jolivet A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
26136718818eSStefano Zampini C->ops->transposematmultsymbolic = MatTransposeMatMultSymbolic_MPIDense_MPIDense;
26146718818eSStefano Zampini C->ops->productsymbolic = MatProductSymbolic_AtB;
26153ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
26164222ddf1SHong Zhang }
26174222ddf1SHong Zhang
MatProductSetFromOptions_MPIDense_ABt(Mat C)2618d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatProductSetFromOptions_MPIDense_ABt(Mat C)
2619d71ae5a4SJacob Faibussowitsch {
26204222ddf1SHong Zhang Mat_Product *product = C->product;
26214222ddf1SHong Zhang const char *algTypes[2] = {"allgatherv", "cyclic"};
26224222ddf1SHong Zhang PetscInt alg, nalg = 2;
26234222ddf1SHong Zhang PetscBool flg = PETSC_FALSE;
26244222ddf1SHong Zhang
26254222ddf1SHong Zhang PetscFunctionBegin;
26264222ddf1SHong Zhang /* Set default algorithm */
26274222ddf1SHong Zhang alg = 0; /* default is allgatherv */
26289566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg, "default", &flg));
2629835f2295SStefano Zampini if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
26304222ddf1SHong Zhang
26314222ddf1SHong Zhang /* Get runtime option */
26324222ddf1SHong Zhang if (product->api_user) {
2633d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatMatTransposeMult", "Mat");
26349566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matmattransmult_mpidense_mpidense_via", "Algorithmic approach", "MatMatTransposeMult", algTypes, nalg, algTypes[alg], &alg, &flg));
2635d0609cedSBarry Smith PetscOptionsEnd();
26364222ddf1SHong Zhang } else {
2637d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_ABt", "Mat");
26389566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_ABt", algTypes, nalg, algTypes[alg], &alg, &flg));
2639d0609cedSBarry Smith PetscOptionsEnd();
26404222ddf1SHong Zhang }
2641835f2295SStefano Zampini if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
26424222ddf1SHong Zhang
26434222ddf1SHong Zhang C->ops->mattransposemultsymbolic = MatMatTransposeMultSymbolic_MPIDense_MPIDense;
26444222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_ABt;
26453ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
26464222ddf1SHong Zhang }
26474222ddf1SHong Zhang
MatProductSetFromOptions_MPIDense(Mat C)26485d8c7819SPierre Jolivet static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat C)
2649d71ae5a4SJacob Faibussowitsch {
26504222ddf1SHong Zhang Mat_Product *product = C->product;
26514222ddf1SHong Zhang
26524222ddf1SHong Zhang PetscFunctionBegin;
26534222ddf1SHong Zhang switch (product->type) {
2654d71ae5a4SJacob Faibussowitsch case MATPRODUCT_AB:
2655d71ae5a4SJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIDense_AB(C));
2656d71ae5a4SJacob Faibussowitsch break;
2657d71ae5a4SJacob Faibussowitsch case MATPRODUCT_AtB:
2658d71ae5a4SJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIDense_AtB(C));
2659d71ae5a4SJacob Faibussowitsch break;
2660d71ae5a4SJacob Faibussowitsch case MATPRODUCT_ABt:
2661d71ae5a4SJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIDense_ABt(C));
2662d71ae5a4SJacob Faibussowitsch break;
2663d71ae5a4SJacob Faibussowitsch default:
2664d71ae5a4SJacob Faibussowitsch break;
26654222ddf1SHong Zhang }
26663ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS);
26674222ddf1SHong Zhang }
2668*d484b384SBoris Martin
MatDenseScatter_Private(PetscSF sf,Mat X,Mat Y,InsertMode mode,ScatterMode smode)2669*d484b384SBoris Martin PetscErrorCode MatDenseScatter_Private(PetscSF sf, Mat X, Mat Y, InsertMode mode, ScatterMode smode)
2670*d484b384SBoris Martin {
2671*d484b384SBoris Martin const PetscScalar *in;
2672*d484b384SBoris Martin PetscScalar *out;
2673*d484b384SBoris Martin PetscSF vsf;
2674*d484b384SBoris Martin PetscInt N, ny, rld, lld;
2675*d484b384SBoris Martin PetscMemType mtype[2];
2676*d484b384SBoris Martin MPI_Op op = MPI_OP_NULL;
2677*d484b384SBoris Martin
2678*d484b384SBoris Martin PetscFunctionBegin;
2679*d484b384SBoris Martin PetscValidHeaderSpecific(sf, PETSCSF_CLASSID, 1);
2680*d484b384SBoris Martin PetscValidHeaderSpecific(X, MAT_CLASSID, 2);
2681*d484b384SBoris Martin PetscValidHeaderSpecific(Y, MAT_CLASSID, 3);
2682*d484b384SBoris Martin if (mode == INSERT_VALUES) op = MPI_REPLACE;
2683*d484b384SBoris Martin else if (mode == ADD_VALUES) op = MPIU_SUM;
2684*d484b384SBoris Martin else if (mode == MAX_VALUES) op = MPIU_MAX;
2685*d484b384SBoris Martin else if (mode == MIN_VALUES) op = MPIU_MIN;
2686*d484b384SBoris Martin PetscCheck(op != MPI_OP_NULL, PetscObjectComm((PetscObject)sf), PETSC_ERR_SUP, "Unsupported InsertMode %d in MatDenseScatter_Private()", mode);
2687*d484b384SBoris Martin PetscCheck(smode == SCATTER_FORWARD || smode == SCATTER_REVERSE, PetscObjectComm((PetscObject)sf), PETSC_ERR_SUP, "Unsupported ScatterMode %d in MatDenseScatter_Private()", smode);
2688*d484b384SBoris Martin PetscCall(MatGetSize(X, NULL, &N));
2689*d484b384SBoris Martin PetscCall(MatGetSize(Y, NULL, &ny));
2690*d484b384SBoris Martin PetscCheck(N == ny, PetscObjectComm((PetscObject)sf), PETSC_ERR_ARG_SIZ, "Matrix column sizes must match: %" PetscInt_FMT " != %" PetscInt_FMT, N, ny);
2691*d484b384SBoris Martin PetscCall(MatDenseGetLDA(X, &rld));
2692*d484b384SBoris Martin PetscCall(MatDenseGetLDA(Y, &lld));
2693*d484b384SBoris Martin /* get cached or create new strided PetscSF when the number of columns is greater than one */
2694*d484b384SBoris Martin if (N > 1) {
2695*d484b384SBoris Martin PetscCall(PetscObjectQuery((PetscObject)sf, "_MatDenseScatter_StridedSF", (PetscObject *)&vsf));
2696*d484b384SBoris Martin if (vsf) {
2697*d484b384SBoris Martin PetscInt nr[2], nl[2];
2698*d484b384SBoris Martin
2699*d484b384SBoris Martin PetscCall(PetscSFGetGraph(sf, nr, nl, NULL, NULL));
2700*d484b384SBoris Martin PetscCall(PetscSFGetGraph(vsf, nr + 1, nl + 1, NULL, NULL));
2701*d484b384SBoris Martin if (N * nr[0] != nr[1] || N * nl[0] != nl[1]) vsf = NULL;
2702*d484b384SBoris Martin }
2703*d484b384SBoris Martin if (!vsf) {
2704*d484b384SBoris Martin PetscCall(PetscSFCreateStridedSF(sf, N, rld, lld, &vsf));
2705*d484b384SBoris Martin PetscCall(PetscObjectCompose((PetscObject)sf, "_MatDenseScatter_StridedSF", (PetscObject)vsf));
2706*d484b384SBoris Martin PetscCall(PetscObjectDereference((PetscObject)vsf));
2707*d484b384SBoris Martin }
2708*d484b384SBoris Martin } else vsf = sf;
2709*d484b384SBoris Martin /* the output array is accessed in read and write mode,
2710*d484b384SBoris Martin but write-only in the INSERT_VALUES case could be worth exploring */
2711*d484b384SBoris Martin PetscCall(MatDenseGetArrayReadAndMemType(X, &in, &mtype[0]));
2712*d484b384SBoris Martin PetscCall(MatDenseGetArrayAndMemType(Y, &out, &mtype[1]));
2713*d484b384SBoris Martin if (smode == SCATTER_FORWARD) {
2714*d484b384SBoris Martin PetscCall(PetscSFBcastWithMemTypeBegin(vsf, vsf->vscat.unit, mtype[0], in, mtype[1], out, op));
2715*d484b384SBoris Martin PetscCall(PetscSFBcastEnd(vsf, vsf->vscat.unit, in, out, op));
2716*d484b384SBoris Martin } else {
2717*d484b384SBoris Martin PetscCall(PetscSFReduceWithMemTypeBegin(vsf, vsf->vscat.unit, mtype[0], in, mtype[1], out, op));
2718*d484b384SBoris Martin PetscCall(PetscSFReduceEnd(vsf, vsf->vscat.unit, in, out, op));
2719*d484b384SBoris Martin }
2720*d484b384SBoris Martin PetscCall(MatDenseRestoreArrayAndMemType(Y, &out));
2721*d484b384SBoris Martin PetscCall(MatDenseRestoreArrayReadAndMemType(X, &in));
2722*d484b384SBoris Martin PetscFunctionReturn(PETSC_SUCCESS);
2723*d484b384SBoris Martin }
2724