Actual source code: mpidense.c
1: /*
2: Basic functions for basic parallel dense matrices.
3: Portions of this code are under:
4: Copyright (c) 2022 Advanced Micro Devices, Inc. All rights reserved.
5: */
7: #include <../src/mat/impls/dense/mpi/mpidense.h>
8: #include <../src/mat/impls/aij/mpi/mpiaij.h>
9: #include <petscblaslapack.h>
10: #include <petsc/private/vecimpl.h>
11: #include <petsc/private/deviceimpl.h>
12: #include <petsc/private/sfimpl.h>
14: /*@
15: MatDenseGetLocalMatrix - For a `MATMPIDENSE` or `MATSEQDENSE` matrix returns the sequential
16: matrix that represents the operator. For sequential matrices it returns itself.
18: Input Parameter:
19: . A - the sequential or MPI `MATDENSE` matrix
21: Output Parameter:
22: . B - the inner matrix
24: Level: intermediate
26: .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MATMPIDENSE`, `MATSEQDENSE`
27: @*/
28: PetscErrorCode MatDenseGetLocalMatrix(Mat A, Mat *B)
29: {
30: Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
31: PetscBool flg;
33: PetscFunctionBegin;
35: PetscAssertPointer(B, 2);
36: PetscCall(PetscObjectBaseTypeCompare((PetscObject)A, MATMPIDENSE, &flg));
37: if (flg) *B = mat->A;
38: else {
39: PetscCall(PetscObjectBaseTypeCompare((PetscObject)A, MATSEQDENSE, &flg));
40: PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "Not for matrix type %s", ((PetscObject)A)->type_name);
41: *B = A;
42: }
43: PetscFunctionReturn(PETSC_SUCCESS);
44: }
46: static PetscErrorCode MatCopy_MPIDense(Mat A, Mat B, MatStructure s)
47: {
48: Mat_MPIDense *Amat = (Mat_MPIDense *)A->data;
49: Mat_MPIDense *Bmat = (Mat_MPIDense *)B->data;
51: PetscFunctionBegin;
52: /* If the two matrices don't have the same copy implementation, they aren't compatible for fast copy. */
53: if (A->ops->copy != B->ops->copy) {
54: PetscCall(MatCopy_Basic(A, B, s));
55: PetscFunctionReturn(PETSC_SUCCESS);
56: }
57: PetscCall(MatCopy(Amat->A, Bmat->A, s));
58: PetscFunctionReturn(PETSC_SUCCESS);
59: }
61: PetscErrorCode MatShift_MPIDense(Mat A, PetscScalar alpha)
62: {
63: Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
64: PetscInt j, lda, rstart = A->rmap->rstart, rend = A->rmap->rend, rend2;
65: PetscScalar *v;
67: PetscFunctionBegin;
68: PetscCall(MatDenseGetArray(mat->A, &v));
69: PetscCall(MatDenseGetLDA(mat->A, &lda));
70: rend2 = PetscMin(rend, A->cmap->N);
71: if (rend2 > rstart) {
72: for (j = rstart; j < rend2; j++) v[j - rstart + j * lda] += alpha;
73: PetscCall(PetscLogFlops(rend2 - rstart));
74: }
75: PetscCall(MatDenseRestoreArray(mat->A, &v));
76: PetscFunctionReturn(PETSC_SUCCESS);
77: }
79: static PetscErrorCode MatGetRow_MPIDense(Mat A, PetscInt row, PetscInt *nz, PetscInt **idx, PetscScalar **v)
80: {
81: Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
82: PetscInt lrow, rstart = A->rmap->rstart, rend = A->rmap->rend;
84: PetscFunctionBegin;
85: PetscCheck(row >= rstart && row < rend, PETSC_COMM_SELF, PETSC_ERR_SUP, "only local rows");
86: lrow = row - rstart;
87: PetscCall(MatGetRow(mat->A, lrow, nz, (const PetscInt **)idx, (const PetscScalar **)v));
88: PetscFunctionReturn(PETSC_SUCCESS);
89: }
91: static PetscErrorCode MatRestoreRow_MPIDense(Mat A, PetscInt row, PetscInt *nz, PetscInt **idx, PetscScalar **v)
92: {
93: Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
94: PetscInt lrow, rstart = A->rmap->rstart, rend = A->rmap->rend;
96: PetscFunctionBegin;
97: PetscCheck(row >= rstart && row < rend, PETSC_COMM_SELF, PETSC_ERR_SUP, "only local rows");
98: lrow = row - rstart;
99: PetscCall(MatRestoreRow(mat->A, lrow, nz, (const PetscInt **)idx, (const PetscScalar **)v));
100: PetscFunctionReturn(PETSC_SUCCESS);
101: }
103: static PetscErrorCode MatGetDiagonalBlock_MPIDense(Mat A, Mat *a)
104: {
105: Mat_MPIDense *mdn = (Mat_MPIDense *)A->data;
106: PetscInt m = A->rmap->n, rstart = A->rmap->rstart;
107: PetscScalar *array;
108: MPI_Comm comm;
109: PetscBool flg;
110: Mat B;
112: PetscFunctionBegin;
113: PetscCall(MatHasCongruentLayouts(A, &flg));
114: PetscCheck(flg, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only square matrices supported.");
115: PetscCall(PetscObjectQuery((PetscObject)A, "DiagonalBlock", (PetscObject *)&B));
116: if (!B) { /* This should use MatDenseGetSubMatrix (not create), but we would need a call like MatRestoreDiagonalBlock */
117: #if PetscDefined(HAVE_CUDA)
118: PetscCall(PetscObjectTypeCompare((PetscObject)mdn->A, MATSEQDENSECUDA, &flg));
119: 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);
120: #elif PetscDefined(HAVE_HIP)
121: PetscCall(PetscObjectTypeCompare((PetscObject)mdn->A, MATSEQDENSEHIP, &flg));
122: 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);
123: #endif
124: PetscCall(PetscObjectGetComm((PetscObject)mdn->A, &comm));
125: PetscCall(MatCreate(comm, &B));
126: PetscCall(MatSetSizes(B, m, m, m, m));
127: PetscCall(MatSetType(B, ((PetscObject)mdn->A)->type_name));
128: PetscCall(MatDenseGetArrayRead(mdn->A, (const PetscScalar **)&array));
129: PetscCall(MatSeqDenseSetPreallocation(B, array + m * rstart));
130: PetscCall(MatDenseRestoreArrayRead(mdn->A, (const PetscScalar **)&array));
131: PetscCall(PetscObjectCompose((PetscObject)A, "DiagonalBlock", (PetscObject)B));
132: *a = B;
133: PetscCall(MatDestroy(&B));
134: } else *a = B;
135: PetscFunctionReturn(PETSC_SUCCESS);
136: }
138: static PetscErrorCode MatSetValues_MPIDense(Mat mat, PetscInt m, const PetscInt idxm[], PetscInt n, const PetscInt idxn[], const PetscScalar v[], InsertMode addv)
139: {
140: Mat_MPIDense *A = (Mat_MPIDense *)mat->data;
141: PetscInt i, j, rstart = mat->rmap->rstart, rend = mat->rmap->rend, row;
142: PetscBool roworiented = A->roworiented;
144: PetscFunctionBegin;
145: for (i = 0; i < m; i++) {
146: if (idxm[i] < 0) continue;
147: PetscCheck(idxm[i] < mat->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large");
148: if (idxm[i] >= rstart && idxm[i] < rend) {
149: row = idxm[i] - rstart;
150: if (roworiented) {
151: PetscCall(MatSetValues(A->A, 1, &row, n, idxn, PetscSafePointerPlusOffset(v, i * n), addv));
152: } else {
153: for (j = 0; j < n; j++) {
154: if (idxn[j] < 0) continue;
155: PetscCheck(idxn[j] < mat->cmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Column too large");
156: PetscCall(MatSetValues(A->A, 1, &row, 1, &idxn[j], PetscSafePointerPlusOffset(v, i + j * m), addv));
157: }
158: }
159: } else if (!A->donotstash) {
160: mat->assembled = PETSC_FALSE;
161: if (roworiented) {
162: PetscCall(MatStashValuesRow_Private(&mat->stash, idxm[i], n, idxn, PetscSafePointerPlusOffset(v, i * n), PETSC_FALSE));
163: } else {
164: PetscCall(MatStashValuesCol_Private(&mat->stash, idxm[i], n, idxn, PetscSafePointerPlusOffset(v, i), m, PETSC_FALSE));
165: }
166: }
167: }
168: PetscFunctionReturn(PETSC_SUCCESS);
169: }
171: static PetscErrorCode MatGetValues_MPIDense(Mat mat, PetscInt m, const PetscInt idxm[], PetscInt n, const PetscInt idxn[], PetscScalar v[])
172: {
173: Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
174: PetscInt i, j, rstart = mat->rmap->rstart, rend = mat->rmap->rend, row;
175: PetscBool roworiented = mdn->roworiented;
176: PetscScalar *value;
178: PetscFunctionBegin;
179: for (i = 0; i < m; i++) {
180: if (idxm[i] < 0) continue; /* negative row */
181: PetscCheck(idxm[i] < mat->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Row too large");
182: PetscCheck(idxm[i] >= rstart && idxm[i] < rend, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only local values currently supported");
183: row = idxm[i] - rstart;
184: for (j = 0; j < n; j++) {
185: if (idxn[j] < 0) continue; /* negative column */
186: PetscCheck(idxn[j] < mat->cmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Column too large");
187: value = roworiented ? &v[j + i * n] : &v[i + j * m];
188: PetscCall(MatGetValues(mdn->A, 1, &row, 1, &idxn[j], value));
189: }
190: }
191: PetscFunctionReturn(PETSC_SUCCESS);
192: }
194: static PetscErrorCode MatDenseGetLDA_MPIDense(Mat A, PetscInt *lda)
195: {
196: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
198: PetscFunctionBegin;
199: PetscCall(MatDenseGetLDA(a->A, lda));
200: PetscFunctionReturn(PETSC_SUCCESS);
201: }
203: static PetscErrorCode MatDenseSetLDA_MPIDense(Mat A, PetscInt lda)
204: {
205: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
206: MatType mtype = MATSEQDENSE;
208: PetscFunctionBegin;
209: if (!a->A) {
210: PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
211: PetscCall(PetscLayoutSetUp(A->rmap));
212: PetscCall(PetscLayoutSetUp(A->cmap));
213: PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
214: PetscCall(MatSetSizes(a->A, A->rmap->n, A->cmap->N, A->rmap->n, A->cmap->N));
215: #if PetscDefined(HAVE_CUDA)
216: PetscBool iscuda;
217: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSECUDA, &iscuda));
218: if (iscuda) mtype = MATSEQDENSECUDA;
219: #elif PetscDefined(HAVE_HIP)
220: PetscBool iship;
221: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSEHIP, &iship));
222: if (iship) mtype = MATSEQDENSEHIP;
223: #endif
224: PetscCall(MatSetType(a->A, mtype));
225: }
226: PetscCall(MatDenseSetLDA(a->A, lda));
227: PetscFunctionReturn(PETSC_SUCCESS);
228: }
230: static PetscErrorCode MatDenseGetArray_MPIDense(Mat A, PetscScalar **array)
231: {
232: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
234: PetscFunctionBegin;
235: PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
236: PetscCall(MatDenseGetArray(a->A, array));
237: PetscFunctionReturn(PETSC_SUCCESS);
238: }
240: static PetscErrorCode MatDenseGetArrayRead_MPIDense(Mat A, PetscScalar **array)
241: {
242: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
244: PetscFunctionBegin;
245: PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
246: PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)array));
247: PetscFunctionReturn(PETSC_SUCCESS);
248: }
250: static PetscErrorCode MatDenseGetArrayWrite_MPIDense(Mat A, PetscScalar **array)
251: {
252: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
254: PetscFunctionBegin;
255: PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
256: PetscCall(MatDenseGetArrayWrite(a->A, array));
257: PetscFunctionReturn(PETSC_SUCCESS);
258: }
260: static PetscErrorCode MatDensePlaceArray_MPIDense(Mat A, const PetscScalar *array)
261: {
262: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
264: PetscFunctionBegin;
265: PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
266: PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
267: PetscCall(MatDensePlaceArray(a->A, array));
268: PetscFunctionReturn(PETSC_SUCCESS);
269: }
271: static PetscErrorCode MatDenseResetArray_MPIDense(Mat A)
272: {
273: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
275: PetscFunctionBegin;
276: PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
277: PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
278: PetscCall(MatDenseResetArray(a->A));
279: PetscFunctionReturn(PETSC_SUCCESS);
280: }
282: static PetscErrorCode MatDenseReplaceArray_MPIDense(Mat A, const PetscScalar *array)
283: {
284: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
286: PetscFunctionBegin;
287: PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
288: PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
289: PetscCall(MatDenseReplaceArray(a->A, array));
290: PetscFunctionReturn(PETSC_SUCCESS);
291: }
293: static PetscErrorCode MatCreateSubMatrix_MPIDense(Mat A, IS isrow, IS iscol, MatReuse scall, Mat *B)
294: {
295: Mat_MPIDense *mat = (Mat_MPIDense *)A->data, *newmatd;
296: PetscInt lda, i, j, rstart, rend, nrows, ncols, Ncols, nlrows, nlcols;
297: const PetscInt *irow, *icol;
298: const PetscScalar *v;
299: PetscScalar *bv;
300: Mat newmat;
301: IS iscol_local;
302: MPI_Comm comm_is, comm_mat;
304: PetscFunctionBegin;
305: PetscCall(PetscObjectGetComm((PetscObject)A, &comm_mat));
306: PetscCall(PetscObjectGetComm((PetscObject)iscol, &comm_is));
307: PetscCheck(comm_mat == comm_is, PETSC_COMM_SELF, PETSC_ERR_ARG_NOTSAMECOMM, "IS communicator must match matrix communicator");
309: PetscCall(ISAllGather(iscol, &iscol_local));
310: PetscCall(ISGetIndices(isrow, &irow));
311: PetscCall(ISGetIndices(iscol_local, &icol));
312: PetscCall(ISGetLocalSize(isrow, &nrows));
313: PetscCall(ISGetLocalSize(iscol, &ncols));
314: PetscCall(ISGetSize(iscol, &Ncols)); /* global number of columns, size of iscol_local */
316: /* No parallel redistribution currently supported! Should really check each index set
317: to confirm that it is OK. ... Currently supports only submatrix same partitioning as
318: original matrix! */
320: PetscCall(MatGetLocalSize(A, &nlrows, &nlcols));
321: PetscCall(MatGetOwnershipRange(A, &rstart, &rend));
323: /* Check submatrix call */
324: if (scall == MAT_REUSE_MATRIX) {
325: /* SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Reused submatrix wrong size"); */
326: /* Really need to test rows and column sizes! */
327: newmat = *B;
328: } else {
329: /* Create and fill new matrix */
330: PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &newmat));
331: PetscCall(MatSetSizes(newmat, nrows, ncols, PETSC_DECIDE, Ncols));
332: PetscCall(MatSetType(newmat, ((PetscObject)A)->type_name));
333: PetscCall(MatMPIDenseSetPreallocation(newmat, NULL));
334: }
336: /* Now extract the data pointers and do the copy, column at a time */
337: newmatd = (Mat_MPIDense *)newmat->data;
338: PetscCall(MatDenseGetArray(newmatd->A, &bv));
339: PetscCall(MatDenseGetArrayRead(mat->A, &v));
340: PetscCall(MatDenseGetLDA(mat->A, &lda));
341: for (i = 0; i < Ncols; i++) {
342: const PetscScalar *av = v + lda * icol[i];
343: for (j = 0; j < nrows; j++) *bv++ = av[irow[j] - rstart];
344: }
345: PetscCall(MatDenseRestoreArrayRead(mat->A, &v));
346: PetscCall(MatDenseRestoreArray(newmatd->A, &bv));
348: /* Assemble the matrices so that the correct flags are set */
349: PetscCall(MatAssemblyBegin(newmat, MAT_FINAL_ASSEMBLY));
350: PetscCall(MatAssemblyEnd(newmat, MAT_FINAL_ASSEMBLY));
352: /* Free work space */
353: PetscCall(ISRestoreIndices(isrow, &irow));
354: PetscCall(ISRestoreIndices(iscol_local, &icol));
355: PetscCall(ISDestroy(&iscol_local));
356: *B = newmat;
357: PetscFunctionReturn(PETSC_SUCCESS);
358: }
360: static PetscErrorCode MatDenseRestoreArray_MPIDense(Mat A, PetscScalar **array)
361: {
362: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
364: PetscFunctionBegin;
365: PetscCall(MatDenseRestoreArray(a->A, array));
366: PetscFunctionReturn(PETSC_SUCCESS);
367: }
369: static PetscErrorCode MatDenseRestoreArrayRead_MPIDense(Mat A, PetscScalar **array)
370: {
371: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
373: PetscFunctionBegin;
374: PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)array));
375: PetscFunctionReturn(PETSC_SUCCESS);
376: }
378: static PetscErrorCode MatDenseRestoreArrayWrite_MPIDense(Mat A, PetscScalar **array)
379: {
380: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
382: PetscFunctionBegin;
383: PetscCall(MatDenseRestoreArrayWrite(a->A, array));
384: PetscFunctionReturn(PETSC_SUCCESS);
385: }
387: static PetscErrorCode MatAssemblyBegin_MPIDense(Mat mat, MatAssemblyType mode)
388: {
389: Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
390: PetscInt nstash, reallocs;
392: PetscFunctionBegin;
393: if (mdn->donotstash || mat->nooffprocentries) PetscFunctionReturn(PETSC_SUCCESS);
395: PetscCall(MatStashScatterBegin_Private(mat, &mat->stash, mat->rmap->range));
396: PetscCall(MatStashGetInfo_Private(&mat->stash, &nstash, &reallocs));
397: PetscCall(PetscInfo(mdn->A, "Stash has %" PetscInt_FMT " entries, uses %" PetscInt_FMT " mallocs.\n", nstash, reallocs));
398: PetscFunctionReturn(PETSC_SUCCESS);
399: }
401: static PetscErrorCode MatAssemblyEnd_MPIDense(Mat mat, MatAssemblyType mode)
402: {
403: Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
404: PetscInt i, *row, *col, flg, j, rstart, ncols;
405: PetscMPIInt n;
406: PetscScalar *val;
408: PetscFunctionBegin;
409: if (!mdn->donotstash && !mat->nooffprocentries) {
410: /* wait on receives */
411: while (1) {
412: PetscCall(MatStashScatterGetMesg_Private(&mat->stash, &n, &row, &col, &val, &flg));
413: if (!flg) break;
415: for (i = 0; i < n;) {
416: /* Now identify the consecutive vals belonging to the same row */
417: for (j = i, rstart = row[j]; j < n; j++) {
418: if (row[j] != rstart) break;
419: }
420: if (j < n) ncols = j - i;
421: else ncols = n - i;
422: /* Now assemble all these values with a single function call */
423: PetscCall(MatSetValues_MPIDense(mat, 1, row + i, ncols, col + i, val + i, mat->insertmode));
424: i = j;
425: }
426: }
427: PetscCall(MatStashScatterEnd_Private(&mat->stash));
428: }
430: PetscCall(MatAssemblyBegin(mdn->A, mode));
431: PetscCall(MatAssemblyEnd(mdn->A, mode));
432: PetscFunctionReturn(PETSC_SUCCESS);
433: }
435: static PetscErrorCode MatZeroEntries_MPIDense(Mat A)
436: {
437: Mat_MPIDense *l = (Mat_MPIDense *)A->data;
439: PetscFunctionBegin;
440: PetscCall(MatZeroEntries(l->A));
441: PetscFunctionReturn(PETSC_SUCCESS);
442: }
444: static PetscErrorCode MatSetInf_MPIDense(Mat A)
445: {
446: Mat_MPIDense *l = (Mat_MPIDense *)A->data;
448: PetscFunctionBegin;
449: PetscCall(MatFlag(l->A, 1));
450: PetscFunctionReturn(PETSC_SUCCESS);
451: }
453: static PetscErrorCode MatZeroRows_MPIDense(Mat A, PetscInt n, const PetscInt rows[], PetscScalar diag, Vec x, Vec b)
454: {
455: Mat_MPIDense *l = (Mat_MPIDense *)A->data;
456: PetscInt i, len, *lrows;
458: PetscFunctionBegin;
459: /* get locally owned rows */
460: PetscCall(PetscLayoutMapLocal(A->rmap, n, rows, &len, &lrows, NULL));
461: /* fix right-hand side if needed */
462: if (x && b) {
463: const PetscScalar *xx;
464: PetscScalar *bb;
466: PetscCall(VecGetArrayRead(x, &xx));
467: PetscCall(VecGetArrayWrite(b, &bb));
468: for (i = 0; i < len; ++i) bb[lrows[i]] = diag * xx[lrows[i]];
469: PetscCall(VecRestoreArrayRead(x, &xx));
470: PetscCall(VecRestoreArrayWrite(b, &bb));
471: }
472: PetscCall(MatZeroRows(l->A, len, lrows, 0.0, NULL, NULL));
473: if (diag != 0.0) {
474: Vec d;
476: PetscCall(MatCreateVecs(A, NULL, &d));
477: PetscCall(VecSet(d, diag));
478: PetscCall(MatDiagonalSet(A, d, INSERT_VALUES));
479: PetscCall(VecDestroy(&d));
480: }
481: PetscCall(PetscFree(lrows));
482: PetscFunctionReturn(PETSC_SUCCESS);
483: }
485: PETSC_INTERN PetscErrorCode MatMult_SeqDense(Mat, Vec, Vec);
486: PETSC_INTERN PetscErrorCode MatMultAdd_SeqDense(Mat, Vec, Vec, Vec);
487: PETSC_INTERN PetscErrorCode MatMultTranspose_SeqDense(Mat, Vec, Vec);
488: PETSC_INTERN PetscErrorCode MatMultTransposeAdd_SeqDense(Mat, Vec, Vec, Vec);
490: static PetscErrorCode MatMultColumnRange_MPIDense(Mat mat, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end)
491: {
492: Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
493: const PetscScalar *ax;
494: PetscScalar *ay;
495: PetscMemType axmtype, aymtype;
497: PetscFunctionBegin;
498: if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
499: PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
500: PetscCall(VecGetArrayWriteAndMemType(mdn->lvec, &ay, &aymtype));
501: PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
502: PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
503: PetscCall(VecRestoreArrayWriteAndMemType(mdn->lvec, &ay));
504: PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
505: PetscUseMethod(mdn->A, "MatMultColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (mdn->A, mdn->lvec, yy, c_start, c_end));
506: PetscFunctionReturn(PETSC_SUCCESS);
507: }
509: static PetscErrorCode MatMult_MPIDense(Mat mat, Vec xx, Vec yy)
510: {
511: Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
512: const PetscScalar *ax;
513: PetscScalar *ay;
514: PetscMemType axmtype, aymtype;
516: PetscFunctionBegin;
517: if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
518: PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
519: PetscCall(VecGetArrayWriteAndMemType(mdn->lvec, &ay, &aymtype));
520: PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
521: PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
522: PetscCall(VecRestoreArrayWriteAndMemType(mdn->lvec, &ay));
523: PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
524: PetscUseTypeMethod(mdn->A, mult, mdn->lvec, yy);
525: PetscFunctionReturn(PETSC_SUCCESS);
526: }
528: static PetscErrorCode MatGetMultPetscSF_MPIDense(Mat A, PetscSF *sf)
529: {
530: Mat_MPIDense *mdn = (Mat_MPIDense *)A->data;
532: PetscFunctionBegin;
533: *sf = mdn->Mvctx;
534: PetscFunctionReturn(PETSC_SUCCESS);
535: }
537: static PetscErrorCode MatMultAddColumnRange_MPIDense(Mat mat, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end)
538: {
539: Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
540: const PetscScalar *ax;
541: PetscScalar *ay;
542: PetscMemType axmtype, aymtype;
544: PetscFunctionBegin;
545: if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
546: PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
547: PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype));
548: PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
549: PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
550: PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay));
551: PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
552: PetscUseMethod(mdn->A, "MatMultAddColumnRange_C", (Mat, Vec, Vec, Vec, PetscInt, PetscInt), (mdn->A, mdn->lvec, yy, zz, c_start, c_end));
553: PetscFunctionReturn(PETSC_SUCCESS);
554: }
556: static PetscErrorCode MatMultAdd_MPIDense(Mat mat, Vec xx, Vec yy, Vec zz)
557: {
558: Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
559: const PetscScalar *ax;
560: PetscScalar *ay;
561: PetscMemType axmtype, aymtype;
563: PetscFunctionBegin;
564: if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(mat));
565: PetscCall(VecGetArrayReadAndMemType(xx, &ax, &axmtype));
566: PetscCall(VecGetArrayAndMemType(mdn->lvec, &ay, &aymtype));
567: PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPI_REPLACE));
568: PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, ax, ay, MPI_REPLACE));
569: PetscCall(VecRestoreArrayAndMemType(mdn->lvec, &ay));
570: PetscCall(VecRestoreArrayReadAndMemType(xx, &ax));
571: PetscUseTypeMethod(mdn->A, multadd, mdn->lvec, yy, zz);
572: PetscFunctionReturn(PETSC_SUCCESS);
573: }
575: static PetscErrorCode MatMultHermitianTransposeColumnRange_MPIDense(Mat A, Vec xx, Vec yy, PetscInt c_start, PetscInt c_end)
576: {
577: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
578: const PetscScalar *ax;
579: PetscScalar *ay;
580: PetscMemType axmtype, aymtype;
581: PetscInt r_start, r_end;
582: PetscInt c_start_local, c_end_local;
584: PetscFunctionBegin;
585: if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
586: PetscCall(VecZeroEntries(a->lvec));
587: PetscCall(VecGetOwnershipRange(yy, &r_start, &r_end));
588: c_start_local = PetscMax(c_start, r_start);
589: c_end_local = PetscMin(c_end, r_end);
590: PetscCall(VecGetArrayAndMemType(yy, &ay, &aymtype));
591: if (c_end_local > c_start_local) {
592: if (PetscMemTypeHost(aymtype)) {
593: PetscCall(PetscArrayzero(&ay[c_start_local], (size_t)(c_end_local - c_start_local)));
594: } else {
595: PetscCall(PetscDeviceRegisterMemory(ay, aymtype, sizeof(*ay) * ((size_t)(r_end - r_start))));
596: PetscCall(PetscDeviceArrayZero(NULL, &ay[c_start_local], (size_t)(c_end_local - c_start_local)));
597: }
598: }
599: PetscUseMethod(a->A, "MatMultHermitianTransposeColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (a->A, xx, a->lvec, c_start, c_end));
600: PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
601: PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
602: PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
603: PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
604: PetscCall(VecRestoreArrayAndMemType(yy, &ay));
605: PetscFunctionReturn(PETSC_SUCCESS);
606: }
608: static PetscErrorCode MatMultTransposeKernel_MPIDense(Mat A, Vec xx, Vec yy, PetscBool herm)
609: {
610: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
611: const PetscScalar *ax;
612: PetscScalar *ay;
613: PetscMemType axmtype, aymtype;
615: PetscFunctionBegin;
616: if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
617: PetscCall(VecSet(yy, 0.0));
618: if (herm) PetscUseTypeMethod(a->A, multhermitiantranspose, xx, a->lvec);
619: else PetscUseTypeMethod(a->A, multtranspose, xx, a->lvec);
620: PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
621: PetscCall(VecGetArrayAndMemType(yy, &ay, &aymtype));
622: PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
623: PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
624: PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
625: PetscCall(VecRestoreArrayAndMemType(yy, &ay));
626: PetscFunctionReturn(PETSC_SUCCESS);
627: }
629: static PetscErrorCode MatMultHermitianTransposeAddColumnRange_MPIDense(Mat A, Vec xx, Vec yy, Vec zz, PetscInt c_start, PetscInt c_end)
630: {
631: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
632: const PetscScalar *ax;
633: PetscScalar *ay;
634: PetscMemType axmtype, aymtype;
636: PetscFunctionBegin;
637: if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
638: PetscCall(VecCopy(yy, zz));
639: PetscCall(VecZeroEntries(a->lvec));
640: PetscUseMethod(a->A, "MatMultHermitianTransposeColumnRange_C", (Mat, Vec, Vec, PetscInt, PetscInt), (a->A, xx, a->lvec, c_start, c_end));
641: PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
642: PetscCall(VecGetArrayAndMemType(zz, &ay, &aymtype));
643: PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
644: PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
645: PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
646: PetscCall(VecRestoreArrayAndMemType(zz, &ay));
647: PetscFunctionReturn(PETSC_SUCCESS);
648: }
650: static PetscErrorCode MatMultTransposeAddKernel_MPIDense(Mat A, Vec xx, Vec yy, Vec zz, PetscBool herm)
651: {
652: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
653: const PetscScalar *ax;
654: PetscScalar *ay;
655: PetscMemType axmtype, aymtype;
657: PetscFunctionBegin;
658: if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
659: PetscCall(VecCopy(yy, zz));
660: if (herm) PetscUseTypeMethod(a->A, multhermitiantranspose, xx, a->lvec);
661: else PetscUseTypeMethod(a->A, multtranspose, xx, a->lvec);
662: PetscCall(VecGetArrayReadAndMemType(a->lvec, &ax, &axmtype));
663: PetscCall(VecGetArrayAndMemType(zz, &ay, &aymtype));
664: PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, MPIU_SCALAR, axmtype, ax, aymtype, ay, MPIU_SUM));
665: PetscCall(PetscSFReduceEnd(a->Mvctx, MPIU_SCALAR, ax, ay, MPIU_SUM));
666: PetscCall(VecRestoreArrayReadAndMemType(a->lvec, &ax));
667: PetscCall(VecRestoreArrayAndMemType(zz, &ay));
668: PetscFunctionReturn(PETSC_SUCCESS);
669: }
671: static PetscErrorCode MatMultTranspose_MPIDense(Mat A, Vec xx, Vec yy)
672: {
673: PetscFunctionBegin;
674: PetscCall(MatMultTransposeKernel_MPIDense(A, xx, yy, PETSC_FALSE));
675: PetscFunctionReturn(PETSC_SUCCESS);
676: }
678: static PetscErrorCode MatMultTransposeAdd_MPIDense(Mat A, Vec xx, Vec yy, Vec zz)
679: {
680: PetscFunctionBegin;
681: PetscCall(MatMultTransposeAddKernel_MPIDense(A, xx, yy, zz, PETSC_FALSE));
682: PetscFunctionReturn(PETSC_SUCCESS);
683: }
685: static PetscErrorCode MatMultHermitianTranspose_MPIDense(Mat A, Vec xx, Vec yy)
686: {
687: PetscFunctionBegin;
688: PetscCall(MatMultTransposeKernel_MPIDense(A, xx, yy, PETSC_TRUE));
689: PetscFunctionReturn(PETSC_SUCCESS);
690: }
692: static PetscErrorCode MatMultHermitianTransposeAdd_MPIDense(Mat A, Vec xx, Vec yy, Vec zz)
693: {
694: PetscFunctionBegin;
695: PetscCall(MatMultTransposeAddKernel_MPIDense(A, xx, yy, zz, PETSC_TRUE));
696: PetscFunctionReturn(PETSC_SUCCESS);
697: }
699: PetscErrorCode MatGetDiagonal_MPIDense(Mat A, Vec v)
700: {
701: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
702: PetscInt lda, len, i, nl, ng, m = A->rmap->n, radd;
703: PetscScalar *x;
704: const PetscScalar *av;
706: PetscFunctionBegin;
707: PetscCall(VecGetArray(v, &x));
708: PetscCall(VecGetSize(v, &ng));
709: PetscCheck(ng == A->rmap->N, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Nonconforming mat and vec");
710: PetscCall(VecGetLocalSize(v, &nl));
711: len = PetscMin(a->A->rmap->n, a->A->cmap->n);
712: radd = A->rmap->rstart * m;
713: PetscCall(MatDenseGetArrayRead(a->A, &av));
714: PetscCall(MatDenseGetLDA(a->A, &lda));
715: for (i = 0; i < len; i++) x[i] = av[radd + i * lda + i];
716: PetscCall(MatDenseRestoreArrayRead(a->A, &av));
717: if (nl - i > 0) PetscCall(PetscArrayzero(x + i, nl - i));
718: PetscCall(VecRestoreArray(v, &x));
719: PetscFunctionReturn(PETSC_SUCCESS);
720: }
722: static PetscErrorCode MatDestroy_MPIDense(Mat mat)
723: {
724: Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
726: PetscFunctionBegin;
727: PetscCall(PetscLogObjectState((PetscObject)mat, "Rows=%" PetscInt_FMT ", Cols=%" PetscInt_FMT, mat->rmap->N, mat->cmap->N));
728: PetscCall(MatStashDestroy_Private(&mat->stash));
729: PetscCheck(!mdn->vecinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
730: PetscCheck(!mdn->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
731: PetscCall(MatDestroy(&mdn->A));
732: PetscCall(VecDestroy(&mdn->lvec));
733: PetscCall(PetscSFDestroy(&mdn->Mvctx));
734: PetscCall(VecDestroy(&mdn->cvec));
735: PetscCall(MatDestroy(&mdn->cmat));
737: PetscCall(PetscFree(mat->data));
738: PetscCall(PetscObjectChangeTypeName((PetscObject)mat, NULL));
740: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", NULL));
741: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", NULL));
742: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", NULL));
743: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", NULL));
744: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", NULL));
745: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", NULL));
746: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", NULL));
747: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", NULL));
748: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", NULL));
749: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", NULL));
750: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", NULL));
751: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", NULL));
752: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", NULL));
753: #if PetscDefined(HAVE_ELEMENTAL)
754: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", NULL));
755: #endif
756: #if PetscDefined(HAVE_SCALAPACK) && (PetscDefined(USE_REAL_SINGLE) || PetscDefined(USE_REAL_DOUBLE))
757: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", NULL));
758: #endif
759: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", NULL));
760: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", NULL));
761: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", NULL));
762: #if PetscDefined(HAVE_CUDA)
763: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", NULL));
764: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", NULL));
765: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", NULL));
766: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidensecuda_mpidense_C", NULL));
767: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidensecuda_C", NULL));
768: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidensecuda_C", NULL));
769: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensecuda_mpiaij_C", NULL));
770: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensecuda_mpiaijcusparse_C", NULL));
771: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArray_C", NULL));
772: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArrayRead_C", NULL));
773: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAGetArrayWrite_C", NULL));
774: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArray_C", NULL));
775: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArrayRead_C", NULL));
776: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDARestoreArrayWrite_C", NULL));
777: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAPlaceArray_C", NULL));
778: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAResetArray_C", NULL));
779: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDAReplaceArray_C", NULL));
780: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseCUDASetPreallocation_C", NULL));
781: #endif
782: #if PetscDefined(HAVE_HIP)
783: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", NULL));
784: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", NULL));
785: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", NULL));
786: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidensehip_mpidense_C", NULL));
787: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidensehip_C", NULL));
788: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidensehip_C", NULL));
789: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensehip_mpiaij_C", NULL));
790: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidensehip_mpiaijhipsparse_C", NULL));
791: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArray_C", NULL));
792: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArrayRead_C", NULL));
793: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPGetArrayWrite_C", NULL));
794: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArray_C", NULL));
795: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArrayRead_C", NULL));
796: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPRestoreArrayWrite_C", NULL));
797: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPPlaceArray_C", NULL));
798: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPResetArray_C", NULL));
799: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPReplaceArray_C", NULL));
800: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseHIPSetPreallocation_C", NULL));
801: #endif
802: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", NULL));
803: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", NULL));
804: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", NULL));
805: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", NULL));
806: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", NULL));
807: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", NULL));
808: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", NULL));
809: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", NULL));
810: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", NULL));
811: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", NULL));
812: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultColumnRange_C", NULL));
813: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", NULL));
814: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", NULL));
815: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", NULL));
816: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatGetMultPetscSF_C", NULL));
817: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseUpdateColumnLayout_C", NULL));
819: PetscCall(PetscObjectCompose((PetscObject)mat, "DiagonalBlock", NULL));
820: PetscFunctionReturn(PETSC_SUCCESS);
821: }
823: #include <petscdraw.h>
824: static PetscErrorCode MatView_MPIDense_ASCIIorDraworSocket(Mat mat, PetscViewer viewer)
825: {
826: Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data;
827: PetscMPIInt rank;
828: PetscViewerType vtype;
829: PetscBool isascii, isdraw;
830: PetscViewer sviewer;
831: PetscViewerFormat format;
833: PetscFunctionBegin;
834: PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)mat), &rank));
835: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
836: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
837: if (isascii) {
838: PetscCall(PetscViewerGetType(viewer, &vtype));
839: PetscCall(PetscViewerGetFormat(viewer, &format));
840: if (format == PETSC_VIEWER_ASCII_INFO_DETAIL) {
841: MatInfo info;
842: PetscCall(MatGetInfo(mat, MAT_LOCAL, &info));
843: PetscCall(PetscViewerASCIIPushSynchronized(viewer));
844: 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,
845: (PetscInt)info.memory));
846: PetscCall(PetscViewerFlush(viewer));
847: PetscCall(PetscViewerASCIIPopSynchronized(viewer));
848: if (mdn->Mvctx) PetscCall(PetscSFView(mdn->Mvctx, viewer));
849: PetscFunctionReturn(PETSC_SUCCESS);
850: } else if (format == PETSC_VIEWER_ASCII_INFO) {
851: PetscFunctionReturn(PETSC_SUCCESS);
852: }
853: } else if (isdraw) {
854: PetscDraw draw;
855: PetscBool isnull;
857: PetscCall(PetscViewerDrawGetDraw(viewer, 0, &draw));
858: PetscCall(PetscDrawIsNull(draw, &isnull));
859: if (isnull) PetscFunctionReturn(PETSC_SUCCESS);
860: }
862: {
863: /* assemble the entire matrix onto first processor. */
864: Mat A;
865: PetscInt M = mat->rmap->N, N = mat->cmap->N, m, row, i, nz;
866: PetscInt *cols;
867: PetscScalar *vals;
869: PetscCall(MatCreate(PetscObjectComm((PetscObject)mat), &A));
870: if (rank == 0) {
871: PetscCall(MatSetSizes(A, M, N, M, N));
872: } else {
873: PetscCall(MatSetSizes(A, 0, 0, M, N));
874: }
875: /* Since this is a temporary matrix, MATMPIDENSE instead of ((PetscObject)A)->type_name here is probably acceptable. */
876: PetscCall(MatSetType(A, MATMPIDENSE));
877: PetscCall(MatMPIDenseSetPreallocation(A, NULL));
879: /* Copy the matrix ... This isn't the most efficient means,
880: but it's quick for now */
881: A->insertmode = INSERT_VALUES;
883: row = mat->rmap->rstart;
884: m = mdn->A->rmap->n;
885: for (i = 0; i < m; i++) {
886: PetscCall(MatGetRow_MPIDense(mat, row, &nz, &cols, &vals));
887: PetscCall(MatSetValues_MPIDense(A, 1, &row, nz, cols, vals, INSERT_VALUES));
888: PetscCall(MatRestoreRow_MPIDense(mat, row, &nz, &cols, &vals));
889: row++;
890: }
892: PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
893: PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
894: PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
895: if (rank == 0) {
896: PetscCall(PetscObjectSetName((PetscObject)((Mat_MPIDense *)A->data)->A, ((PetscObject)mat)->name));
897: PetscCall(MatView_SeqDense(((Mat_MPIDense *)A->data)->A, sviewer));
898: }
899: PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
900: PetscCall(MatDestroy(&A));
901: }
902: PetscFunctionReturn(PETSC_SUCCESS);
903: }
905: static PetscErrorCode MatView_MPIDense(Mat mat, PetscViewer viewer)
906: {
907: PetscBool isascii, isbinary, isdraw, issocket;
909: PetscFunctionBegin;
910: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
911: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
912: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERSOCKET, &issocket));
913: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERDRAW, &isdraw));
915: if (isascii || issocket || isdraw) PetscCall(MatView_MPIDense_ASCIIorDraworSocket(mat, viewer));
916: else if (isbinary) PetscCall(MatView_Dense_Binary(mat, viewer));
917: PetscFunctionReturn(PETSC_SUCCESS);
918: }
920: static PetscErrorCode MatGetInfo_MPIDense(Mat A, MatInfoType flag, MatInfo *info)
921: {
922: Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
923: Mat mdn = mat->A;
924: PetscLogDouble irecv[5];
926: PetscFunctionBegin;
927: info->block_size = 1.0;
929: PetscCall(MatGetInfo(mdn, MAT_LOCAL, info));
931: irecv[0] = info->nz_used;
932: irecv[1] = info->nz_allocated;
933: irecv[2] = info->nz_unneeded;
934: irecv[3] = info->memory;
935: irecv[4] = info->mallocs;
936: if (flag == MAT_LOCAL) {
937: info->nz_used = irecv[0];
938: info->nz_allocated = irecv[1];
939: info->nz_unneeded = irecv[2];
940: info->memory = irecv[3];
941: info->mallocs = irecv[4];
942: } else if (flag == MAT_GLOBAL_MAX) {
943: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_MAX, PetscObjectComm((PetscObject)A)));
945: info->nz_used = irecv[0];
946: info->nz_allocated = irecv[1];
947: info->nz_unneeded = irecv[2];
948: info->memory = irecv[3];
949: info->mallocs = irecv[4];
950: } else if (flag == MAT_GLOBAL_SUM) {
951: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, irecv, 5, MPIU_PETSCLOGDOUBLE, MPI_SUM, PetscObjectComm((PetscObject)A)));
953: info->nz_used = irecv[0];
954: info->nz_allocated = irecv[1];
955: info->nz_unneeded = irecv[2];
956: info->memory = irecv[3];
957: info->mallocs = irecv[4];
958: }
959: info->fill_ratio_given = 0; /* no parallel LU/ILU/Cholesky */
960: info->fill_ratio_needed = 0;
961: info->factor_mallocs = 0;
962: PetscFunctionReturn(PETSC_SUCCESS);
963: }
965: static PetscErrorCode MatSetOption_MPIDense(Mat A, MatOption op, PetscBool flg)
966: {
967: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
969: PetscFunctionBegin;
970: switch (op) {
971: case MAT_NEW_NONZERO_LOCATIONS:
972: case MAT_NEW_NONZERO_LOCATION_ERR:
973: case MAT_NEW_NONZERO_ALLOCATION_ERR:
974: MatCheckPreallocated(A, 1);
975: PetscCall(MatSetOption(a->A, op, flg));
976: break;
977: case MAT_ROW_ORIENTED:
978: MatCheckPreallocated(A, 1);
979: a->roworiented = flg;
980: PetscCall(MatSetOption(a->A, op, flg));
981: break;
982: case MAT_IGNORE_OFF_PROC_ENTRIES:
983: a->donotstash = flg;
984: break;
985: case MAT_SYMMETRIC:
986: case MAT_STRUCTURALLY_SYMMETRIC:
987: case MAT_HERMITIAN:
988: case MAT_SYMMETRY_ETERNAL:
989: case MAT_STRUCTURAL_SYMMETRY_ETERNAL:
990: case MAT_SPD:
991: case MAT_SPD_ETERNAL:
992: /* if the diagonal matrix is square it inherits some of the properties above */
993: if (a->A && A->rmap->n == A->cmap->n) PetscCall(MatSetOption(a->A, op, flg));
994: break;
995: default:
996: break;
997: }
998: PetscFunctionReturn(PETSC_SUCCESS);
999: }
1001: static PetscErrorCode MatDiagonalScale_MPIDense(Mat A, Vec ll, Vec rr)
1002: {
1003: Mat_MPIDense *mdn = (Mat_MPIDense *)A->data;
1004: PetscInt s1, s2, s3;
1006: PetscFunctionBegin;
1007: PetscCall(MatGetLocalSize(A, &s2, &s3));
1008: if (ll) {
1009: PetscCall(VecGetLocalSize(ll, &s1));
1010: PetscCheck(s1 == s2, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Left scaling vector non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT, s1, s2);
1011: }
1012: if (rr) {
1013: PetscCall(VecGetLocalSize(rr, &s1));
1014: PetscCheck(s1 == s3, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Right scaling vector non-conforming local size, %" PetscInt_FMT " != %" PetscInt_FMT, s1, s3);
1015: /* gather the right scaling into the local column layout, staying on the device when the Vecs live there */
1016: if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
1017: PetscCall(VecScatterBegin(mdn->Mvctx, rr, mdn->lvec, INSERT_VALUES, SCATTER_FORWARD));
1018: PetscCall(VecScatterEnd(mdn->Mvctx, rr, mdn->lvec, INSERT_VALUES, SCATTER_FORWARD));
1019: }
1020: /* the local matrix holds exactly the local rows and all the columns, so the local part of ll applies to it directly;
1021: the operation is called rather than MatDiagonalScale() because ll is the parallel Vec and the interface checks that
1022: it shares the communicator of the sequential matrix, as in MatDiagonalScale_MPIAIJ() */
1023: PetscUseTypeMethod(mdn->A, diagonalscale, ll, rr ? mdn->lvec : NULL);
1024: PetscCall(PetscObjectStateIncrease((PetscObject)mdn->A));
1025: PetscFunctionReturn(PETSC_SUCCESS);
1026: }
1028: static PetscErrorCode MatNorm_MPIDense(Mat A, NormType type, PetscReal *nrm)
1029: {
1030: Mat_MPIDense *mdn = (Mat_MPIDense *)A->data;
1031: PetscInt i, j, lda;
1032: PetscMPIInt size;
1033: const PetscScalar *av;
1035: PetscFunctionBegin;
1036: PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)A), &size));
1037: if (size == 1) {
1038: PetscCall(MatNorm(mdn->A, type, nrm));
1039: } else {
1040: if (type == NORM_FROBENIUS) {
1041: PetscCall(MatNorm(mdn->A, NORM_FROBENIUS, nrm));
1042: *nrm *= *nrm;
1043: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, nrm, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
1044: *nrm = PetscSqrtReal(*nrm);
1045: } else if (type == NORM_1) {
1046: PetscReal *tmp;
1048: PetscCall(PetscCalloc1(A->cmap->N, &tmp));
1049: *nrm = 0.0;
1050: PetscCall(MatDenseGetArrayRead(mdn->A, &av));
1051: PetscCall(MatDenseGetLDA(mdn->A, &lda));
1052: for (j = 0; j < mdn->A->cmap->n; j++) {
1053: for (i = 0; i < mdn->A->rmap->n; i++) tmp[j] += PetscAbsScalar(av[i + j * lda]);
1054: }
1055: PetscCall(MatDenseRestoreArrayRead(mdn->A, &av));
1056: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, tmp, A->cmap->N, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)A)));
1057: for (j = 0; j < A->cmap->N; j++) {
1058: if (tmp[j] > *nrm) *nrm = tmp[j];
1059: }
1060: PetscCall(PetscFree(tmp));
1061: PetscCall(PetscLogFlops(A->cmap->n * A->rmap->n));
1062: } else if (type == NORM_INFINITY) { /* max row norm */
1063: PetscCall(MatNorm(mdn->A, type, nrm));
1064: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, nrm, 1, MPIU_REAL, MPIU_MAX, PetscObjectComm((PetscObject)A)));
1065: } else SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_SUP, "Unsupported norm type %s", NormTypes[type]);
1066: }
1067: PetscFunctionReturn(PETSC_SUCCESS);
1068: }
1070: static PetscErrorCode MatTranspose_MPIDense(Mat A, MatReuse reuse, Mat *matout)
1071: {
1072: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1073: Mat B;
1074: PetscInt M = A->rmap->N, N = A->cmap->N, m, n, *rwork, rstart = A->rmap->rstart;
1075: PetscInt j, i, lda;
1076: PetscScalar *v;
1078: PetscFunctionBegin;
1079: if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(A, *matout));
1080: if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_INPLACE_MATRIX) {
1081: PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &B));
1082: PetscCall(MatSetSizes(B, A->cmap->n, A->rmap->n, N, M));
1083: PetscCall(MatSetType(B, ((PetscObject)A)->type_name));
1084: PetscCall(MatMPIDenseSetPreallocation(B, NULL));
1085: } else B = *matout;
1087: m = a->A->rmap->n;
1088: n = a->A->cmap->n;
1089: PetscCall(MatDenseGetArrayRead(a->A, (const PetscScalar **)&v));
1090: PetscCall(MatDenseGetLDA(a->A, &lda));
1091: PetscCall(PetscMalloc1(m, &rwork));
1092: for (i = 0; i < m; i++) rwork[i] = rstart + i;
1093: for (j = 0; j < n; j++) {
1094: PetscCall(MatSetValues(B, 1, &j, m, rwork, v, INSERT_VALUES));
1095: v = PetscSafePointerPlusOffset(v, lda);
1096: }
1097: PetscCall(MatDenseRestoreArrayRead(a->A, (const PetscScalar **)&v));
1098: PetscCall(PetscFree(rwork));
1099: PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
1100: PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
1101: if (reuse == MAT_INITIAL_MATRIX || reuse == MAT_REUSE_MATRIX) {
1102: *matout = B;
1103: } else {
1104: PetscCall(MatHeaderMerge(A, &B));
1105: }
1106: PetscFunctionReturn(PETSC_SUCCESS);
1107: }
1109: static PetscErrorCode MatDuplicate_MPIDense(Mat, MatDuplicateOption, Mat *);
1110: PETSC_INTERN PetscErrorCode MatScale_MPIDense(Mat, PetscScalar);
1112: static PetscErrorCode MatSetUp_MPIDense(Mat A)
1113: {
1114: PetscFunctionBegin;
1115: PetscCall(PetscLayoutSetUp(A->rmap));
1116: PetscCall(PetscLayoutSetUp(A->cmap));
1117: if (!A->preallocated) PetscCall(MatMPIDenseSetPreallocation(A, NULL));
1118: PetscFunctionReturn(PETSC_SUCCESS);
1119: }
1121: static PetscErrorCode MatAXPY_MPIDense(Mat Y, PetscScalar alpha, Mat X, MatStructure str)
1122: {
1123: Mat_MPIDense *A = (Mat_MPIDense *)Y->data, *B = (Mat_MPIDense *)X->data;
1125: PetscFunctionBegin;
1126: PetscCall(MatAXPY(A->A, alpha, B->A, str));
1127: PetscFunctionReturn(PETSC_SUCCESS);
1128: }
1130: static PetscErrorCode MatConjugate_MPIDense(Mat mat)
1131: {
1132: Mat_MPIDense *a = (Mat_MPIDense *)mat->data;
1134: PetscFunctionBegin;
1135: PetscCall(MatConjugate(a->A));
1136: PetscFunctionReturn(PETSC_SUCCESS);
1137: }
1139: static PetscErrorCode MatRealPart_MPIDense(Mat A)
1140: {
1141: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1143: PetscFunctionBegin;
1144: PetscCall(MatRealPart(a->A));
1145: PetscFunctionReturn(PETSC_SUCCESS);
1146: }
1148: static PetscErrorCode MatImaginaryPart_MPIDense(Mat A)
1149: {
1150: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1152: PetscFunctionBegin;
1153: PetscCall(MatImaginaryPart(a->A));
1154: PetscFunctionReturn(PETSC_SUCCESS);
1155: }
1157: static PetscErrorCode MatGetColumnVector_MPIDense(Mat A, Vec v, PetscInt col)
1158: {
1159: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1161: PetscFunctionBegin;
1162: PetscCheck(a->A, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing local matrix");
1163: PetscCheck(a->A->ops->getcolumnvector, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Missing get column operation");
1164: PetscUseTypeMethod(a->A, getcolumnvector, v, col);
1165: PetscFunctionReturn(PETSC_SUCCESS);
1166: }
1168: PETSC_INTERN PetscErrorCode MatGetColumnReductions_SeqDense(Mat, PetscInt, PetscReal *);
1170: static PetscErrorCode MatGetColumnReductions_MPIDense(Mat A, PetscInt type, PetscReal *reductions)
1171: {
1172: PetscInt i, m, n;
1173: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1175: PetscFunctionBegin;
1176: PetscCall(MatGetSize(A, &m, &n));
1177: if (type == REDUCTION_MEAN_REALPART) {
1178: PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_REALPART, reductions));
1179: } else if (type == REDUCTION_MEAN_IMAGINARYPART) {
1180: PetscCall(MatGetColumnReductions_SeqDense(a->A, (PetscInt)REDUCTION_SUM_IMAGINARYPART, reductions));
1181: } else {
1182: PetscCall(MatGetColumnReductions_SeqDense(a->A, type, reductions));
1183: }
1184: if (type == NORM_2) {
1185: for (i = 0; i < n; i++) reductions[i] *= reductions[i];
1186: }
1187: if (type == NORM_INFINITY) {
1188: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, reductions, n, MPIU_REAL, MPIU_MAX, A->hdr.comm));
1189: } else {
1190: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, reductions, n, MPIU_REAL, MPIU_SUM, A->hdr.comm));
1191: }
1192: if (type == NORM_2) {
1193: for (i = 0; i < n; i++) reductions[i] = PetscSqrtReal(reductions[i]);
1194: } else if (type == REDUCTION_MEAN_REALPART || type == REDUCTION_MEAN_IMAGINARYPART) {
1195: for (i = 0; i < n; i++) reductions[i] /= m;
1196: }
1197: PetscFunctionReturn(PETSC_SUCCESS);
1198: }
1200: static PetscErrorCode MatSetRandom_MPIDense(Mat x, PetscRandom rctx)
1201: {
1202: Mat_MPIDense *d = (Mat_MPIDense *)x->data;
1204: PetscFunctionBegin;
1205: PetscCall(MatSetRandom(d->A, rctx));
1206: #if PetscDefined(HAVE_DEVICE)
1207: x->offloadmask = d->A->offloadmask;
1208: #endif
1209: PetscCall(PetscObjectStateIncrease((PetscObject)x));
1210: PetscFunctionReturn(PETSC_SUCCESS);
1211: }
1213: static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat);
1214: static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat);
1215: static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat, Mat, PetscReal, Mat);
1216: static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat, Mat, Mat);
1217: static PetscErrorCode MatEqual_MPIDense(Mat, Mat, PetscBool *);
1218: static PetscErrorCode MatLoad_MPIDense(Mat, PetscViewer);
1219: static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat);
1221: static struct _MatOps MatOps_Values = {MatSetValues_MPIDense,
1222: MatGetRow_MPIDense,
1223: MatRestoreRow_MPIDense,
1224: MatMult_MPIDense,
1225: /* 4*/ MatMultAdd_MPIDense,
1226: MatMultTranspose_MPIDense,
1227: MatMultTransposeAdd_MPIDense,
1228: NULL,
1229: NULL,
1230: NULL,
1231: /* 10*/ NULL,
1232: NULL,
1233: NULL,
1234: NULL,
1235: MatTranspose_MPIDense,
1236: /* 15*/ MatGetInfo_MPIDense,
1237: MatEqual_MPIDense,
1238: MatGetDiagonal_MPIDense,
1239: MatDiagonalScale_MPIDense,
1240: MatNorm_MPIDense,
1241: /* 20*/ MatAssemblyBegin_MPIDense,
1242: MatAssemblyEnd_MPIDense,
1243: MatSetOption_MPIDense,
1244: MatZeroEntries_MPIDense,
1245: /* 24*/ MatZeroRows_MPIDense,
1246: NULL,
1247: NULL,
1248: NULL,
1249: NULL,
1250: /* 29*/ MatSetUp_MPIDense,
1251: NULL,
1252: NULL,
1253: MatGetDiagonalBlock_MPIDense,
1254: MatSetInf_MPIDense,
1255: /* 34*/ MatDuplicate_MPIDense,
1256: NULL,
1257: NULL,
1258: NULL,
1259: NULL,
1260: /* 39*/ MatAXPY_MPIDense,
1261: MatCreateSubMatrices_MPIDense,
1262: NULL,
1263: MatGetValues_MPIDense,
1264: MatCopy_MPIDense,
1265: /* 44*/ NULL,
1266: MatScale_MPIDense,
1267: MatShift_MPIDense,
1268: NULL,
1269: NULL,
1270: /* 49*/ MatSetRandom_MPIDense,
1271: NULL,
1272: NULL,
1273: NULL,
1274: NULL,
1275: /* 54*/ NULL,
1276: NULL,
1277: NULL,
1278: NULL,
1279: NULL,
1280: /* 59*/ MatCreateSubMatrix_MPIDense,
1281: MatDestroy_MPIDense,
1282: MatView_MPIDense,
1283: NULL,
1284: NULL,
1285: /* 64*/ NULL,
1286: NULL,
1287: NULL,
1288: NULL,
1289: NULL,
1290: /* 69*/ NULL,
1291: NULL,
1292: NULL,
1293: NULL,
1294: NULL,
1295: /* 74*/ NULL,
1296: NULL,
1297: NULL,
1298: NULL,
1299: MatLoad_MPIDense,
1300: /* 79*/ NULL,
1301: NULL,
1302: NULL,
1303: NULL,
1304: /* 83*/ NULL,
1305: NULL,
1306: NULL,
1307: NULL,
1308: MatMatTransposeMultSymbolic_MPIDense_MPIDense,
1309: MatMatTransposeMultNumeric_MPIDense_MPIDense,
1310: /* 89*/ NULL,
1311: MatProductSetFromOptions_MPIDense,
1312: NULL,
1313: NULL,
1314: MatConjugate_MPIDense,
1315: /* 94*/ NULL,
1316: NULL,
1317: MatRealPart_MPIDense,
1318: MatImaginaryPart_MPIDense,
1319: NULL,
1320: /*99*/ NULL,
1321: NULL,
1322: NULL,
1323: NULL,
1324: MatGetColumnVector_MPIDense,
1325: /*104*/ NULL,
1326: NULL,
1327: NULL,
1328: NULL,
1329: NULL,
1330: /*109*/ NULL,
1331: NULL,
1332: MatMultHermitianTranspose_MPIDense,
1333: MatMultHermitianTransposeAdd_MPIDense,
1334: NULL,
1335: /*114*/ NULL,
1336: MatGetColumnReductions_MPIDense,
1337: NULL,
1338: NULL,
1339: NULL,
1340: /*120*/ MatTransposeMatMultSymbolic_MPIDense_MPIDense,
1341: MatTransposeMatMultNumeric_MPIDense_MPIDense,
1342: NULL,
1343: NULL,
1344: /*124*/ NULL,
1345: NULL,
1346: NULL,
1347: NULL,
1348: NULL,
1349: /*129*/ NULL,
1350: MatCreateMPIMatConcatenateSeqMat_MPIDense,
1351: NULL,
1352: NULL,
1353: NULL,
1354: /*134*/ NULL,
1355: NULL,
1356: NULL,
1357: NULL,
1358: NULL,
1359: /*139*/ NULL,
1360: NULL,
1361: NULL,
1362: NULL,
1363: NULL,
1364: NULL,
1365: /*144*/ NULL,
1366: NULL,
1367: NULL,
1368: NULL};
1370: static PetscErrorCode MatMPIDenseSetPreallocation_MPIDense(Mat mat, PetscScalar *data)
1371: {
1372: Mat_MPIDense *a = (Mat_MPIDense *)mat->data;
1373: MatType mtype = MATSEQDENSE;
1375: PetscFunctionBegin;
1376: PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)mat), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1377: PetscCall(PetscLayoutSetUp(mat->rmap));
1378: PetscCall(PetscLayoutSetUp(mat->cmap));
1379: if (!a->A) {
1380: PetscCall(MatCreate(PETSC_COMM_SELF, &a->A));
1381: PetscCall(MatSetSizes(a->A, mat->rmap->n, mat->cmap->N, mat->rmap->n, mat->cmap->N));
1382: }
1383: #if PetscDefined(HAVE_CUDA)
1384: PetscBool iscuda;
1385: PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSECUDA, &iscuda));
1386: if (iscuda) mtype = MATSEQDENSECUDA;
1387: #endif
1388: #if PetscDefined(HAVE_HIP)
1389: PetscBool iship;
1390: PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSEHIP, &iship));
1391: if (iship) mtype = MATSEQDENSEHIP;
1392: #endif
1393: PetscCall(MatSetType(a->A, mtype));
1394: PetscCall(MatSeqDenseSetPreallocation(a->A, data));
1395: #if PetscDefined(HAVE_CUDA) || PetscDefined(HAVE_HIP)
1396: mat->offloadmask = a->A->offloadmask;
1397: #endif
1398: mat->preallocated = PETSC_TRUE;
1399: mat->assembled = PETSC_TRUE;
1400: PetscFunctionReturn(PETSC_SUCCESS);
1401: }
1403: PETSC_INTERN PetscErrorCode MatConvert_MPIAIJ_MPIDense(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1404: {
1405: Mat B, C;
1407: PetscFunctionBegin;
1408: PetscCall(MatMPIAIJGetLocalMat(A, MAT_INITIAL_MATRIX, &C));
1409: PetscCall(MatConvert_SeqAIJ_SeqDense(C, MATSEQDENSE, MAT_INITIAL_MATRIX, &B));
1410: PetscCall(MatDestroy(&C));
1411: if (reuse == MAT_REUSE_MATRIX) {
1412: C = *newmat;
1413: } else C = NULL;
1414: PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C));
1415: PetscCall(MatDestroy(&B));
1416: if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &C));
1417: else if (reuse == MAT_INITIAL_MATRIX) *newmat = C;
1418: PetscFunctionReturn(PETSC_SUCCESS);
1419: }
1421: static PetscErrorCode MatConvert_MPIDense_MPIAIJ(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1422: {
1423: Mat B, C;
1425: PetscFunctionBegin;
1426: PetscCall(MatDenseGetLocalMatrix(A, &C));
1427: PetscCall(MatConvert_SeqDense_SeqAIJ(C, MATSEQAIJ, MAT_INITIAL_MATRIX, &B));
1428: if (reuse == MAT_REUSE_MATRIX) {
1429: C = *newmat;
1430: } else C = NULL;
1431: PetscCall(MatCreateMPIMatConcatenateSeqMat(PetscObjectComm((PetscObject)A), B, A->cmap->n, !C ? MAT_INITIAL_MATRIX : MAT_REUSE_MATRIX, &C));
1432: PetscCall(MatDestroy(&B));
1433: if (reuse == MAT_INPLACE_MATRIX) PetscCall(MatHeaderReplace(A, &C));
1434: else if (reuse == MAT_INITIAL_MATRIX) *newmat = C;
1435: PetscFunctionReturn(PETSC_SUCCESS);
1436: }
1438: #if PetscDefined(HAVE_ELEMENTAL)
1439: PETSC_INTERN PetscErrorCode MatConvert_MPIDense_Elemental(Mat A, MatType newtype, MatReuse reuse, Mat *newmat)
1440: {
1441: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1442: Mat mat_elemental;
1443: PetscScalar *v;
1444: PetscInt m = A->rmap->n, N = A->cmap->N, rstart = A->rmap->rstart, i, *rows, *cols, lda;
1446: PetscFunctionBegin;
1447: if (reuse == MAT_REUSE_MATRIX) {
1448: mat_elemental = *newmat;
1449: PetscCall(MatZeroEntries(*newmat));
1450: } else {
1451: PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat_elemental));
1452: PetscCall(MatSetSizes(mat_elemental, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N));
1453: PetscCall(MatSetType(mat_elemental, MATELEMENTAL));
1454: PetscCall(MatSetUp(mat_elemental));
1455: PetscCall(MatSetOption(mat_elemental, MAT_ROW_ORIENTED, PETSC_FALSE));
1456: }
1458: PetscCall(PetscMalloc2(m, &rows, N, &cols));
1459: for (i = 0; i < N; i++) cols[i] = i;
1460: for (i = 0; i < m; i++) rows[i] = rstart + i;
1462: /* PETSc-Elemental interface uses axpy for setting off-processor entries, only ADD_VALUES is allowed */
1463: PetscCall(MatDenseGetArray(A, &v));
1464: PetscCall(MatDenseGetLDA(a->A, &lda));
1465: if (lda == m) PetscCall(MatSetValues(mat_elemental, m, rows, N, cols, v, ADD_VALUES));
1466: else {
1467: for (i = 0; i < N; i++) PetscCall(MatSetValues(mat_elemental, m, rows, 1, &i, v + lda * i, ADD_VALUES));
1468: }
1469: PetscCall(MatAssemblyBegin(mat_elemental, MAT_FINAL_ASSEMBLY));
1470: PetscCall(MatAssemblyEnd(mat_elemental, MAT_FINAL_ASSEMBLY));
1471: PetscCall(MatDenseRestoreArray(A, &v));
1472: PetscCall(PetscFree2(rows, cols));
1474: if (reuse == MAT_INPLACE_MATRIX) {
1475: PetscCall(MatHeaderReplace(A, &mat_elemental));
1476: } else {
1477: *newmat = mat_elemental;
1478: }
1479: PetscFunctionReturn(PETSC_SUCCESS);
1480: }
1481: #endif
1483: static PetscErrorCode MatDenseGetColumn_MPIDense(Mat A, PetscInt col, PetscScalar **vals)
1484: {
1485: Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
1487: PetscFunctionBegin;
1488: PetscCall(MatDenseGetColumn(mat->A, col, vals));
1489: PetscFunctionReturn(PETSC_SUCCESS);
1490: }
1492: static PetscErrorCode MatDenseRestoreColumn_MPIDense(Mat A, PetscScalar **vals)
1493: {
1494: Mat_MPIDense *mat = (Mat_MPIDense *)A->data;
1496: PetscFunctionBegin;
1497: PetscCall(MatDenseRestoreColumn(mat->A, vals));
1498: PetscFunctionReturn(PETSC_SUCCESS);
1499: }
1501: PetscErrorCode MatCreateMPIMatConcatenateSeqMat_MPIDense(MPI_Comm comm, Mat inmat, PetscInt n, MatReuse scall, Mat *outmat)
1502: {
1503: Mat_MPIDense *mat;
1504: PetscInt m, nloc, N;
1506: PetscFunctionBegin;
1507: PetscCall(MatGetSize(inmat, &m, &N));
1508: PetscCall(MatGetLocalSize(inmat, NULL, &nloc));
1509: if (scall == MAT_INITIAL_MATRIX) { /* symbolic phase */
1510: PetscInt sum;
1512: if (n == PETSC_DECIDE) PetscCall(PetscSplitOwnership(comm, &n, &N));
1513: /* Check sum(n) = N */
1514: PetscCallMPI(MPIU_Allreduce(&n, &sum, 1, MPIU_INT, MPI_SUM, comm));
1515: PetscCheck(sum == N, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Sum of local columns %" PetscInt_FMT " != global columns %" PetscInt_FMT, sum, N);
1517: PetscCall(MatCreateDense(comm, m, n, PETSC_DETERMINE, N, NULL, outmat));
1518: PetscCall(MatSetOption(*outmat, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
1519: }
1521: /* numeric phase */
1522: mat = (Mat_MPIDense *)(*outmat)->data;
1523: PetscCall(MatCopy(inmat, mat->A, SAME_NONZERO_PATTERN));
1524: PetscFunctionReturn(PETSC_SUCCESS);
1525: }
1527: PetscErrorCode MatDenseGetColumnVec_MPIDense(Mat A, PetscInt col, Vec *v)
1528: {
1529: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1530: PetscInt lda;
1532: PetscFunctionBegin;
1533: PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1534: PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1535: if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
1536: a->vecinuse = col + 1;
1537: PetscCall(MatDenseGetLDA(a->A, &lda));
1538: PetscCall(MatDenseGetArray(a->A, (PetscScalar **)&a->ptrinuse));
1539: PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
1540: *v = a->cvec;
1541: PetscFunctionReturn(PETSC_SUCCESS);
1542: }
1544: PetscErrorCode MatDenseRestoreColumnVec_MPIDense(Mat A, PetscInt col, Vec *v)
1545: {
1546: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1548: PetscFunctionBegin;
1549: PetscCheck(a->vecinuse, PETSC_COMM_SELF, PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
1550: PetscCheck(a->cvec, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Missing internal column vector");
1551: VecCheckAssembled(a->cvec);
1552: a->vecinuse = 0;
1553: PetscCall(MatDenseRestoreArray(a->A, (PetscScalar **)&a->ptrinuse));
1554: PetscCall(VecResetArray(a->cvec));
1555: if (v) *v = NULL;
1556: PetscFunctionReturn(PETSC_SUCCESS);
1557: }
1559: PetscErrorCode MatDenseGetColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v)
1560: {
1561: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1562: PetscInt lda;
1564: PetscFunctionBegin;
1565: PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1566: PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1567: if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
1568: a->vecinuse = col + 1;
1569: PetscCall(MatDenseGetLDA(a->A, &lda));
1570: PetscCall(MatDenseGetArrayRead(a->A, &a->ptrinuse));
1571: PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
1572: PetscCall(VecLockReadPush(a->cvec));
1573: *v = a->cvec;
1574: PetscFunctionReturn(PETSC_SUCCESS);
1575: }
1577: PetscErrorCode MatDenseRestoreColumnVecRead_MPIDense(Mat A, PetscInt col, Vec *v)
1578: {
1579: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1581: PetscFunctionBegin;
1582: PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
1583: PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector");
1584: VecCheckAssembled(a->cvec);
1585: a->vecinuse = 0;
1586: PetscCall(MatDenseRestoreArrayRead(a->A, &a->ptrinuse));
1587: PetscCall(VecLockReadPop(a->cvec));
1588: PetscCall(VecResetArray(a->cvec));
1589: if (v) *v = NULL;
1590: PetscFunctionReturn(PETSC_SUCCESS);
1591: }
1593: PetscErrorCode MatDenseGetColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v)
1594: {
1595: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1596: PetscInt lda;
1598: PetscFunctionBegin;
1599: PetscCheck(!a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1600: PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1601: if (!a->cvec) PetscCall(MatDenseCreateColumnVec_Private(A, &a->cvec));
1602: a->vecinuse = col + 1;
1603: PetscCall(MatDenseGetLDA(a->A, &lda));
1604: PetscCall(MatDenseGetArrayWrite(a->A, (PetscScalar **)&a->ptrinuse));
1605: PetscCall(VecPlaceArray(a->cvec, PetscSafePointerPlusOffset(a->ptrinuse, (size_t)col * (size_t)lda)));
1606: *v = a->cvec;
1607: PetscFunctionReturn(PETSC_SUCCESS);
1608: }
1610: PetscErrorCode MatDenseRestoreColumnVecWrite_MPIDense(Mat A, PetscInt col, Vec *v)
1611: {
1612: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1614: PetscFunctionBegin;
1615: PetscCheck(a->vecinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetColumnVec() first");
1616: PetscCheck(a->cvec, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal column vector");
1617: VecCheckAssembled(a->cvec);
1618: a->vecinuse = 0;
1619: PetscCall(MatDenseRestoreArrayWrite(a->A, (PetscScalar **)&a->ptrinuse));
1620: PetscCall(VecResetArray(a->cvec));
1621: if (v) *v = NULL;
1622: PetscFunctionReturn(PETSC_SUCCESS);
1623: }
1625: static PetscErrorCode MatDenseGetSubMatrix_MPIDense(Mat A, PetscInt rbegin, PetscInt rend, PetscInt cbegin, PetscInt cend, Mat *v)
1626: {
1627: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1628: Mat_MPIDense *c;
1629: MPI_Comm comm;
1630: PetscInt prbegin, prend, pcbegin, pcend;
1632: PetscFunctionBegin;
1633: PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
1634: PetscCheck(!a->vecinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreColumnVec() first");
1635: PetscCheck(!a->matinuse, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1636: prbegin = PetscMax(0, PetscMin(A->rmap->rend, rbegin) - A->rmap->rstart);
1637: prend = PetscMin(A->rmap->n, PetscMax(0, rend - A->rmap->rstart));
1638: pcbegin = PetscMax(0, PetscMin(A->cmap->rend, cbegin) - A->cmap->rstart);
1639: pcend = PetscMin(A->cmap->n, PetscMax(0, cend - A->cmap->rstart));
1640: if (!a->cmat) {
1641: PetscCall(MatCreate(comm, &a->cmat));
1642: PetscCall(MatSetType(a->cmat, ((PetscObject)A)->type_name));
1643: PetscCall(MatSetVecType(a->cmat, A->defaultvectype));
1644: if (rend - rbegin == A->rmap->N) PetscCall(PetscLayoutReference(A->rmap, &a->cmat->rmap));
1645: else {
1646: PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, prend - prbegin));
1647: PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin));
1648: PetscCall(PetscLayoutSetUp(a->cmat->rmap));
1649: }
1650: if (cend - cbegin == A->cmap->N) PetscCall(PetscLayoutReference(A->cmap, &a->cmat->cmap));
1651: else {
1652: PetscCall(PetscLayoutSetLocalSize(a->cmat->cmap, pcend - pcbegin));
1653: PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin));
1654: PetscCall(PetscLayoutSetUp(a->cmat->cmap));
1655: }
1656: c = (Mat_MPIDense *)a->cmat->data;
1657: c->sub_rbegin = rbegin;
1658: c->sub_rend = rend;
1659: c->sub_cbegin = cbegin;
1660: c->sub_cend = cend;
1661: }
1662: c = (Mat_MPIDense *)a->cmat->data;
1663: if (c->sub_rbegin != rbegin || c->sub_rend != rend) {
1664: PetscCall(PetscLayoutDestroy(&a->cmat->rmap));
1665: PetscCall(PetscLayoutCreate(comm, &a->cmat->rmap));
1666: PetscCall(PetscLayoutSetLocalSize(a->cmat->rmap, prend - prbegin));
1667: PetscCall(PetscLayoutSetSize(a->cmat->rmap, rend - rbegin));
1668: PetscCall(PetscLayoutSetUp(a->cmat->rmap));
1669: c->sub_rbegin = rbegin;
1670: c->sub_rend = rend;
1671: }
1672: if (c->sub_cbegin != cbegin || c->sub_cend != cend) {
1673: // special optimization: check if all columns are owned by rank 0, in which case no communication is necessary
1674: if (cend - cbegin != a->cmat->cmap->N || A->cmap->range[1] != A->cmap->N) {
1675: PetscCall(PetscLayoutDestroy(&a->cmat->cmap));
1676: PetscCall(PetscLayoutCreate(comm, &a->cmat->cmap));
1677: PetscCall(PetscLayoutSetLocalSize(a->cmat->cmap, pcend - pcbegin));
1678: PetscCall(PetscLayoutSetSize(a->cmat->cmap, cend - cbegin));
1679: PetscCall(PetscLayoutSetUp(a->cmat->cmap));
1680: PetscCall(VecDestroy(&c->lvec));
1681: PetscCall(PetscSFDestroy(&c->Mvctx));
1682: }
1683: c->sub_cbegin = cbegin;
1684: c->sub_cend = cend;
1685: }
1686: PetscCheck(!c->A, comm, PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1687: PetscCall(MatDenseGetSubMatrix(a->A, prbegin, prend, cbegin, cend, &c->A));
1689: a->cmat->preallocated = PETSC_TRUE;
1690: a->cmat->assembled = PETSC_TRUE;
1691: #if PetscDefined(HAVE_DEVICE)
1692: a->cmat->offloadmask = c->A->offloadmask;
1693: #endif
1694: a->matinuse = cbegin + 1;
1695: *v = a->cmat;
1696: PetscFunctionReturn(PETSC_SUCCESS);
1697: }
1699: static PetscErrorCode MatDenseRestoreSubMatrix_MPIDense(Mat A, Mat *v)
1700: {
1701: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1702: Mat_MPIDense *c;
1704: PetscFunctionBegin;
1705: PetscCheck(a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseGetSubMatrix() first");
1706: PetscCheck(a->cmat, PetscObjectComm((PetscObject)A), PETSC_ERR_PLIB, "Missing internal matrix");
1707: PetscCheck(*v == a->cmat, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Not the matrix obtained from MatDenseGetSubMatrix()");
1708: a->matinuse = 0;
1709: c = (Mat_MPIDense *)a->cmat->data;
1710: PetscCall(MatDenseRestoreSubMatrix(a->A, &c->A));
1711: *v = NULL;
1712: #if PetscDefined(HAVE_DEVICE)
1713: A->offloadmask = a->A->offloadmask;
1714: #endif
1715: PetscFunctionReturn(PETSC_SUCCESS);
1716: }
1718: static PetscErrorCode MatDenseUpdateColumnLayout_MPIDense(Mat A, PetscLayout clayout)
1719: {
1720: Mat_MPIDense *a = (Mat_MPIDense *)A->data;
1722: PetscFunctionBegin;
1723: if (A->cmap == clayout) PetscFunctionReturn(PETSC_SUCCESS);
1724: PetscCheck(!a->matinuse, PetscObjectComm((PetscObject)A), PETSC_ERR_ORDER, "Need to call MatDenseRestoreSubMatrix() first");
1725: PetscCall(PetscLayoutReference(clayout, &A->cmap));
1726: PetscCall(MatDestroy(&a->cmat));
1727: PetscCall(PetscSFDestroy(&a->Mvctx));
1728: PetscFunctionReturn(PETSC_SUCCESS);
1729: }
1731: /*MC
1732: MATMPIDENSE - MATMPIDENSE = "mpidense" - A matrix type to be used for distributed dense matrices.
1734: Options Database Key:
1735: . -mat_type mpidense - sets the matrix type to `MATMPIDENSE` during a call to `MatSetFromOptions()`
1737: Level: beginner
1739: .seealso: [](ch_matrices), `Mat`, `MatCreateDense()`, `MATSEQDENSE`, `MATDENSE`
1740: M*/
1741: PetscErrorCode MatCreate_MPIDense(Mat mat)
1742: {
1743: Mat_MPIDense *a;
1745: PetscFunctionBegin;
1746: PetscCall(PetscNew(&a));
1747: mat->data = (void *)a;
1748: mat->ops[0] = MatOps_Values;
1750: mat->insertmode = NOT_SET_VALUES;
1752: /* build cache for off array entries formed */
1753: a->donotstash = PETSC_FALSE;
1755: PetscCall(MatStashCreate_Private(PetscObjectComm((PetscObject)mat), 1, &mat->stash));
1757: /* stuff used for matrix vector multiply */
1758: a->lvec = NULL;
1759: a->Mvctx = NULL;
1760: a->roworiented = PETSC_TRUE;
1762: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetLDA_C", MatDenseGetLDA_MPIDense));
1763: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseSetLDA_C", MatDenseSetLDA_MPIDense));
1764: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArray_C", MatDenseGetArray_MPIDense));
1765: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArray_C", MatDenseRestoreArray_MPIDense));
1766: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayRead_C", MatDenseGetArrayRead_MPIDense));
1767: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayRead_C", MatDenseRestoreArrayRead_MPIDense));
1768: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetArrayWrite_C", MatDenseGetArrayWrite_MPIDense));
1769: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreArrayWrite_C", MatDenseRestoreArrayWrite_MPIDense));
1770: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDensePlaceArray_C", MatDensePlaceArray_MPIDense));
1771: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseResetArray_C", MatDenseResetArray_MPIDense));
1772: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseReplaceArray_C", MatDenseReplaceArray_MPIDense));
1773: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVec_C", MatDenseGetColumnVec_MPIDense));
1774: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVec_C", MatDenseRestoreColumnVec_MPIDense));
1775: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecRead_C", MatDenseGetColumnVecRead_MPIDense));
1776: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecRead_C", MatDenseRestoreColumnVecRead_MPIDense));
1777: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumnVecWrite_C", MatDenseGetColumnVecWrite_MPIDense));
1778: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumnVecWrite_C", MatDenseRestoreColumnVecWrite_MPIDense));
1779: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetSubMatrix_C", MatDenseGetSubMatrix_MPIDense));
1780: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreSubMatrix_C", MatDenseRestoreSubMatrix_MPIDense));
1781: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpiaij_mpidense_C", MatConvert_MPIAIJ_MPIDense));
1782: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpiaij_C", MatConvert_MPIDense_MPIAIJ));
1783: #if PetscDefined(HAVE_ELEMENTAL)
1784: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_elemental_C", MatConvert_MPIDense_Elemental));
1785: #endif
1786: #if PetscDefined(HAVE_SCALAPACK) && (PetscDefined(USE_REAL_SINGLE) || PetscDefined(USE_REAL_DOUBLE))
1787: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_scalapack_C", MatConvert_Dense_ScaLAPACK));
1788: #endif
1789: #if PetscDefined(HAVE_CUDA)
1790: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensecuda_C", MatConvert_MPIDense_MPIDenseCUDA));
1791: #endif
1792: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMPIDenseSetPreallocation_C", MatMPIDenseSetPreallocation_MPIDense));
1793: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaij_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
1794: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaij_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
1795: #if PetscDefined(HAVE_CUDA)
1796: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijcusparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
1797: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijcusparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
1798: #endif
1799: #if PetscDefined(HAVE_HIP)
1800: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatConvert_mpidense_mpidensehip_C", MatConvert_MPIDense_MPIDenseHIP));
1801: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpiaijhipsparse_mpidense_C", MatProductSetFromOptions_MPIAIJ_MPIDense));
1802: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatProductSetFromOptions_mpidense_mpiaijhipsparse_C", MatProductSetFromOptions_MPIDense_MPIAIJ));
1803: #endif
1804: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseGetColumn_C", MatDenseGetColumn_MPIDense));
1805: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseRestoreColumn_C", MatDenseRestoreColumn_MPIDense));
1806: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultColumnRange_C", MatMultColumnRange_MPIDense));
1807: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultAddColumnRange_C", MatMultAddColumnRange_MPIDense));
1808: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeColumnRange_C", MatMultHermitianTransposeColumnRange_MPIDense));
1809: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatMultHermitianTransposeAddColumnRange_C", MatMultHermitianTransposeAddColumnRange_MPIDense));
1810: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatGetMultPetscSF_C", MatGetMultPetscSF_MPIDense));
1811: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatDenseUpdateColumnLayout_C", MatDenseUpdateColumnLayout_MPIDense));
1812: PetscCall(PetscObjectChangeTypeName((PetscObject)mat, MATMPIDENSE));
1813: PetscFunctionReturn(PETSC_SUCCESS);
1814: }
1816: /*MC
1817: MATDENSE - MATDENSE = "dense" - A matrix type to be used for dense matrices.
1819: This matrix type is identical to `MATSEQDENSE` when constructed with a single process communicator,
1820: and `MATMPIDENSE` otherwise.
1822: Options Database Key:
1823: . -mat_type dense - sets the matrix type to `MATDENSE` during a call to `MatSetFromOptions()`
1825: Level: beginner
1827: .seealso: [](ch_matrices), `Mat`, `MATSEQDENSE`, `MATMPIDENSE`, `MATDENSECUDA`, `MATDENSEHIP`
1828: M*/
1830: /*@
1831: MatMPIDenseSetPreallocation - Sets the array used to store the matrix entries
1833: Collective
1835: Input Parameters:
1836: + B - the matrix
1837: - data - optional location of matrix data. Set to `NULL` for PETSc
1838: to control all matrix memory allocation.
1840: Level: intermediate
1842: Notes:
1843: The dense format is fully compatible with standard Fortran
1844: storage by columns.
1846: The data input variable is intended primarily for Fortran programmers
1847: who wish to allocate their own matrix memory space. Most users should
1848: set `data` to `NULL`.
1850: .seealso: [](ch_matrices), `Mat`, `MATMPIDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()`
1851: @*/
1852: PetscErrorCode MatMPIDenseSetPreallocation(Mat B, PetscScalar *data)
1853: {
1854: PetscFunctionBegin;
1856: PetscTryMethod(B, "MatMPIDenseSetPreallocation_C", (Mat, PetscScalar *), (B, data));
1857: PetscFunctionReturn(PETSC_SUCCESS);
1858: }
1860: /*@
1861: MatDensePlaceArray - Allows one to replace the array in a `MATDENSE` matrix with an
1862: array provided by the user. This is useful to avoid copying an array
1863: into a matrix
1865: Not Collective
1867: Input Parameters:
1868: + mat - the matrix
1869: - array - the array in column major order
1871: Level: developer
1873: Note:
1874: Adding `const` to `array` was an oversight, see notes in `VecPlaceArray()`.
1876: You can return to the original array with a call to `MatDenseResetArray()`. The user is responsible for freeing this array; it will not be
1877: freed when the matrix is destroyed.
1879: .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDenseResetArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`,
1880: `MatDenseReplaceArray()`
1881: @*/
1882: PetscErrorCode MatDensePlaceArray(Mat mat, const PetscScalar *array)
1883: {
1884: PetscFunctionBegin;
1886: PetscUseMethod(mat, "MatDensePlaceArray_C", (Mat, const PetscScalar *), (mat, array));
1887: PetscCall(PetscObjectStateIncrease((PetscObject)mat));
1888: #if PetscDefined(HAVE_CUDA) || PetscDefined(HAVE_HIP)
1889: mat->offloadmask = PETSC_OFFLOAD_CPU;
1890: #endif
1891: PetscFunctionReturn(PETSC_SUCCESS);
1892: }
1894: /*@
1895: MatDenseResetArray - Resets the matrix array to that it previously had before the call to `MatDensePlaceArray()`
1897: Not Collective
1899: Input Parameter:
1900: . mat - the matrix
1902: Level: developer
1904: Note:
1905: You can only call this after a call to `MatDensePlaceArray()`
1907: .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatDenseGetArray()`, `MatDensePlaceArray()`, `VecPlaceArray()`, `VecGetArray()`, `VecRestoreArray()`, `VecReplaceArray()`, `VecResetArray()`
1908: @*/
1909: PetscErrorCode MatDenseResetArray(Mat mat)
1910: {
1911: PetscFunctionBegin;
1913: PetscUseMethod(mat, "MatDenseResetArray_C", (Mat), (mat));
1914: PetscCall(PetscObjectStateIncrease((PetscObject)mat));
1915: PetscFunctionReturn(PETSC_SUCCESS);
1916: }
1918: /*@
1919: MatDenseReplaceArray - Allows one to replace the array in a dense matrix with an
1920: array provided by the user. This is useful to avoid copying an array
1921: into a matrix
1923: Not Collective
1925: Input Parameters:
1926: + mat - the matrix
1927: - array - the array in column major order
1929: Level: developer
1931: Note:
1932: Adding `const` to `array` was an oversight, see notes in `VecPlaceArray()`.
1934: The memory passed in MUST be obtained with `PetscMalloc()` and CANNOT be
1935: freed by the user. It will be freed when the matrix is destroyed.
1937: .seealso: [](ch_matrices), `Mat`, `MatDensePlaceArray()`, `MatDenseGetArray()`, `VecReplaceArray()`
1938: @*/
1939: PetscErrorCode MatDenseReplaceArray(Mat mat, const PetscScalar *array)
1940: {
1941: PetscFunctionBegin;
1943: PetscUseMethod(mat, "MatDenseReplaceArray_C", (Mat, const PetscScalar *), (mat, array));
1944: PetscCall(PetscObjectStateIncrease((PetscObject)mat));
1945: #if PetscDefined(HAVE_CUDA) || PetscDefined(HAVE_HIP)
1946: mat->offloadmask = PETSC_OFFLOAD_CPU;
1947: #endif
1948: PetscFunctionReturn(PETSC_SUCCESS);
1949: }
1951: /*@
1952: MatCreateDense - Creates a matrix in `MATDENSE` format.
1954: Collective
1956: Input Parameters:
1957: + comm - MPI communicator
1958: . m - number of local rows (or `PETSC_DECIDE` to have calculated if `M` is given)
1959: . n - number of local columns (or `PETSC_DECIDE` to have calculated if `N` is given)
1960: . M - number of global rows (or `PETSC_DECIDE` to have calculated if `m` is given)
1961: . N - number of global columns (or `PETSC_DECIDE` to have calculated if `n` is given)
1962: - data - optional location of matrix data. Set data to `NULL` (`PETSC_NULL_SCALAR_ARRAY` for Fortran users) for PETSc
1963: to control all matrix memory allocation.
1965: Output Parameter:
1966: . A - the matrix
1968: Level: intermediate
1970: Notes:
1971: The dense format is fully compatible with standard Fortran
1972: storage by columns.
1974: Although local portions of the matrix are stored in column-major
1975: order, the matrix is partitioned across MPI ranks by row.
1977: The data input variable is intended primarily for Fortran programmers
1978: who wish to allocate their own matrix memory space. Most users should
1979: set `data` to `NULL` (`PETSC_NULL_SCALAR_ARRAY` for Fortran users).
1981: The user MUST specify either the local or global matrix dimensions
1982: (possibly both).
1984: .seealso: [](ch_matrices), `Mat`, `MATDENSE`, `MatCreate()`, `MatCreateSeqDense()`, `MatSetValues()`
1985: @*/
1986: PetscErrorCode MatCreateDense(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt M, PetscInt N, PetscScalar data[], Mat *A)
1987: {
1988: PetscFunctionBegin;
1989: PetscCall(MatCreate(comm, A));
1990: PetscCall(MatSetSizes(*A, m, n, M, N));
1991: PetscCall(MatSetType(*A, MATDENSE));
1992: PetscCall(MatSeqDenseSetPreallocation(*A, data));
1993: PetscCall(MatMPIDenseSetPreallocation(*A, data));
1994: PetscFunctionReturn(PETSC_SUCCESS);
1995: }
1997: static PetscErrorCode MatDuplicate_MPIDense(Mat A, MatDuplicateOption cpvalues, Mat *newmat)
1998: {
1999: Mat mat;
2000: Mat_MPIDense *a, *oldmat = (Mat_MPIDense *)A->data;
2002: PetscFunctionBegin;
2003: *newmat = NULL;
2004: PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &mat));
2005: PetscCall(MatSetSizes(mat, A->rmap->n, A->cmap->n, A->rmap->N, A->cmap->N));
2006: PetscCall(MatSetType(mat, ((PetscObject)A)->type_name));
2007: a = (Mat_MPIDense *)mat->data;
2009: mat->factortype = A->factortype;
2010: mat->assembled = PETSC_TRUE;
2011: mat->preallocated = PETSC_TRUE;
2013: mat->insertmode = NOT_SET_VALUES;
2014: a->donotstash = oldmat->donotstash;
2016: PetscCall(PetscLayoutReference(A->rmap, &mat->rmap));
2017: PetscCall(PetscLayoutReference(A->cmap, &mat->cmap));
2019: PetscCall(MatDuplicate(oldmat->A, cpvalues, &a->A));
2021: *newmat = mat;
2022: PetscFunctionReturn(PETSC_SUCCESS);
2023: }
2025: static PetscErrorCode MatLoad_MPIDense(Mat newMat, PetscViewer viewer)
2026: {
2027: PetscBool isbinary;
2028: #if PetscDefined(HAVE_HDF5)
2029: PetscBool ishdf5;
2030: #endif
2032: PetscFunctionBegin;
2035: /* force binary viewer to load .info file if it has not yet done so */
2036: PetscCall(PetscViewerSetUp(viewer));
2037: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERBINARY, &isbinary));
2038: #if PetscDefined(HAVE_HDF5)
2039: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERHDF5, &ishdf5));
2040: #endif
2041: if (isbinary) {
2042: PetscCall(MatLoad_Dense_Binary(newMat, viewer));
2043: #if PetscDefined(HAVE_HDF5)
2044: } else if (ishdf5) {
2045: PetscCall(MatLoad_Dense_HDF5(newMat, viewer));
2046: #endif
2047: } 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);
2048: PetscFunctionReturn(PETSC_SUCCESS);
2049: }
2051: static PetscErrorCode MatEqual_MPIDense(Mat A, Mat B, PetscBool *flag)
2052: {
2053: Mat_MPIDense *matB = (Mat_MPIDense *)B->data, *matA = (Mat_MPIDense *)A->data;
2054: Mat a, b;
2056: PetscFunctionBegin;
2057: a = matA->A;
2058: b = matB->A;
2059: PetscCall(MatEqual(a, b, flag));
2060: PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, flag, 1, MPI_C_BOOL, MPI_LAND, PetscObjectComm((PetscObject)A)));
2061: PetscFunctionReturn(PETSC_SUCCESS);
2062: }
2064: static PetscErrorCode MatProductCtxDestroy_MatTransMatMult_MPIDense_MPIDense(PetscCtxRt data)
2065: {
2066: MatProductCtx_TransMatMultDense *atb = *(MatProductCtx_TransMatMultDense **)data;
2068: PetscFunctionBegin;
2069: PetscCall(PetscFree2(atb->sendbuf, atb->recvcounts));
2070: PetscCall(MatDestroy(&atb->atb));
2071: PetscCall(PetscFree(atb));
2072: PetscFunctionReturn(PETSC_SUCCESS);
2073: }
2075: static PetscErrorCode MatProductCtxDestroy_MatMatTransMult_MPIDense_MPIDense(PetscCtxRt data)
2076: {
2077: MatProductCtx_MatTransMultDense *abt = *(MatProductCtx_MatTransMultDense **)data;
2079: PetscFunctionBegin;
2080: PetscCall(PetscFree2(abt->buf[0], abt->buf[1]));
2081: PetscCall(PetscFree2(abt->recvcounts, abt->recvdispls));
2082: PetscCall(PetscFree(abt));
2083: PetscFunctionReturn(PETSC_SUCCESS);
2084: }
2086: static PetscErrorCode MatTransposeMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2087: {
2088: Mat_MPIDense *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
2089: MatProductCtx_TransMatMultDense *atb;
2090: MPI_Comm comm;
2091: PetscMPIInt size, *recvcounts;
2092: PetscScalar *carray, *sendbuf;
2093: const PetscScalar *atbarray;
2094: PetscInt i, cN = C->cmap->N, proc, k, j, lda;
2095: const PetscInt *ranges;
2097: PetscFunctionBegin;
2098: MatCheckProduct(C, 3);
2099: PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2100: atb = (MatProductCtx_TransMatMultDense *)C->product->data;
2101: recvcounts = atb->recvcounts;
2102: sendbuf = atb->sendbuf;
2104: PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
2105: PetscCallMPI(MPI_Comm_size(comm, &size));
2107: /* compute atbarray = aseq^T * bseq */
2108: PetscCall(MatTransposeMatMult(a->A, b->A, atb->atb ? MAT_REUSE_MATRIX : MAT_INITIAL_MATRIX, PETSC_DETERMINE, &atb->atb));
2110: PetscCall(MatGetOwnershipRanges(C, &ranges));
2112: if (ranges[1] == C->rmap->N) {
2113: /* all of the values are being reduced to rank 0: optimize this case to use MPI_Reduce and GPU aware MPI if available */
2114: PetscInt atb_lda, c_lda;
2115: Mat atb_local = atb->atb;
2116: Mat atb_alloc = NULL;
2117: Mat c_local = c->A;
2118: Mat c_alloc = NULL;
2119: PetscMemType atb_memtype, c_memtype;
2120: const PetscScalar *atb_array = NULL;
2121: MPI_Datatype vector_type;
2122: PetscScalar *c_array = NULL;
2123: PetscMPIInt rank;
2125: PetscCallMPI(MPI_Comm_rank(comm, &rank));
2127: PetscCall(MatDenseGetLDA(atb_local, &atb_lda));
2128: if (atb_lda != C->rmap->N) {
2129: // copy atb to a matrix that will have lda == the number of rows
2130: PetscCall(MatDuplicate(atb_local, MAT_DO_NOT_COPY_VALUES, &atb_alloc));
2131: PetscCall(MatCopy(atb_local, atb_alloc, DIFFERENT_NONZERO_PATTERN));
2132: atb_local = atb_alloc;
2133: }
2135: if (rank == 0) {
2136: PetscCall(MatDenseGetLDA(c_local, &c_lda));
2137: if (c_lda != C->rmap->N) {
2138: // copy c to a matrix that will have lda == the number of rows
2139: PetscCall(MatDuplicate(c_local, MAT_DO_NOT_COPY_VALUES, &c_alloc));
2140: c_local = c_alloc;
2141: }
2142: PetscCall(MatZeroEntries(c_local));
2143: }
2144: /* atb_local and c_local have nrows = lda = A->cmap->N and ncols =
2145: * B->cmap->N: use the a->Mvctx to use the best reduction method */
2146: if (!a->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A));
2147: vector_type = MPIU_SCALAR;
2148: if (B->cmap->N > 1) {
2149: PetscMPIInt mpi_N;
2151: PetscCall(PetscMPIIntCast(B->cmap->N, &mpi_N));
2152: PetscCallMPI(MPI_Type_contiguous(mpi_N, MPIU_SCALAR, &vector_type));
2153: PetscCallMPI(MPI_Type_commit(&vector_type));
2154: }
2155: PetscCall(MatDenseGetArrayReadAndMemType(atb_local, &atb_array, &atb_memtype));
2156: PetscCall(MatDenseGetArrayWriteAndMemType(c_local, &c_array, &c_memtype));
2157: PetscCall(PetscSFReduceWithMemTypeBegin(a->Mvctx, vector_type, atb_memtype, atb_array, c_memtype, c_array, MPIU_SUM));
2158: PetscCall(PetscSFReduceEnd(a->Mvctx, vector_type, atb_array, c_array, MPIU_SUM));
2159: PetscCall(MatDenseRestoreArrayWriteAndMemType(c_local, &c_array));
2160: PetscCall(MatDenseRestoreArrayReadAndMemType(atb_local, &atb_array));
2161: if (rank == 0 && c_local != c->A) PetscCall(MatCopy(c_local, c->A, DIFFERENT_NONZERO_PATTERN));
2162: if (B->cmap->N > 1) PetscCallMPI(MPI_Type_free(&vector_type));
2163: PetscCall(MatDestroy(&atb_alloc));
2164: PetscCall(MatDestroy(&c_alloc));
2165: PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2166: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2167: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2168: PetscFunctionReturn(PETSC_SUCCESS);
2169: }
2171: /* arrange atbarray into sendbuf */
2172: PetscCall(MatDenseGetArrayRead(atb->atb, &atbarray));
2173: PetscCall(MatDenseGetLDA(atb->atb, &lda));
2174: for (proc = 0, k = 0; proc < size; proc++) {
2175: for (j = 0; j < cN; j++) {
2176: for (i = ranges[proc]; i < ranges[proc + 1]; i++) sendbuf[k++] = atbarray[i + j * lda];
2177: }
2178: }
2179: PetscCall(MatDenseRestoreArrayRead(atb->atb, &atbarray));
2181: /* sum all atbarray to local values of C */
2182: PetscCall(MatDenseGetArrayWrite(c->A, &carray));
2183: PetscCallMPI(MPI_Reduce_scatter(sendbuf, carray, recvcounts, MPIU_SCALAR, MPIU_SUM, comm));
2184: PetscCall(MatDenseRestoreArrayWrite(c->A, &carray));
2185: PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2186: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2187: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2188: PetscFunctionReturn(PETSC_SUCCESS);
2189: }
2191: static PetscErrorCode MatTransposeMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2192: {
2193: MPI_Comm comm;
2194: PetscMPIInt size;
2195: PetscInt cm = A->cmap->n, cM, cN = B->cmap->N;
2196: MatProductCtx_TransMatMultDense *atb;
2197: PetscBool cisdense = PETSC_FALSE;
2198: const PetscInt *ranges;
2200: PetscFunctionBegin;
2201: MatCheckProduct(C, 4);
2202: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
2203: PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
2204: 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,
2205: A->rmap->rend, B->rmap->rstart, B->rmap->rend);
2207: /* create matrix product C */
2208: PetscCall(MatSetSizes(C, cm, B->cmap->n, A->cmap->N, B->cmap->N));
2209: #if PetscDefined(HAVE_CUDA)
2210: PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSECUDA, ""));
2211: #endif
2212: #if PetscDefined(HAVE_HIP)
2213: PetscCall(PetscObjectTypeCompareAny((PetscObject)C, &cisdense, MATMPIDENSE, MATMPIDENSEHIP, ""));
2214: #endif
2215: if (!cisdense) {
2216: PetscCall(MatSetType(C, ((PetscObject)A)->type_name));
2217: PetscCall(MatSetVecType(C, A->defaultvectype));
2218: }
2219: PetscCall(MatSetUp(C));
2221: /* create data structure for reuse C */
2222: PetscCallMPI(MPI_Comm_size(comm, &size));
2223: PetscCall(PetscNew(&atb));
2224: cM = C->rmap->N;
2225: PetscCall(PetscMalloc2(cM * cN, &atb->sendbuf, size, &atb->recvcounts));
2226: PetscCall(MatGetOwnershipRanges(C, &ranges));
2227: for (PetscMPIInt i = 0; i < size; i++) PetscCall(PetscMPIIntCast((ranges[i + 1] - ranges[i]) * cN, &atb->recvcounts[i]));
2228: C->product->data = atb;
2229: C->product->destroy = MatProductCtxDestroy_MatTransMatMult_MPIDense_MPIDense;
2230: PetscFunctionReturn(PETSC_SUCCESS);
2231: }
2233: static PetscErrorCode MatMatTransposeMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2234: {
2235: MPI_Comm comm;
2236: PetscMPIInt i, size;
2237: PetscInt maxRows, bufsiz;
2238: PetscMPIInt tag;
2239: PetscInt alg;
2240: MatProductCtx_MatTransMultDense *abt;
2241: Mat_Product *product = C->product;
2242: PetscBool flg;
2244: PetscFunctionBegin;
2245: MatCheckProduct(C, 4);
2246: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
2247: /* check local size of A and B */
2248: 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);
2250: PetscCall(PetscStrcmp(product->alg, "allgatherv", &flg));
2251: alg = flg ? 0 : 1;
2253: /* setup matrix product C */
2254: PetscCall(MatSetSizes(C, A->rmap->n, B->rmap->n, A->rmap->N, B->rmap->N));
2255: PetscCall(MatSetType(C, MATMPIDENSE));
2256: /* keep the VecType of A, e.g. VECKOKKOS from MatCreateDenseFromVecType(), but only if A is a MATMPIDENSE like C, since the VecType of a MATMPIDENSECUDA or MATMPIDENSEHIP A, e.g. VECCUDA, is not valid for C */
2257: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSE, &flg));
2258: if (flg) PetscCall(MatSetVecType(C, A->defaultvectype));
2259: PetscCall(MatSetUp(C));
2260: PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag));
2262: /* create data structure for reuse C */
2263: PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
2264: PetscCallMPI(MPI_Comm_size(comm, &size));
2265: PetscCall(PetscNew(&abt));
2266: abt->tag = tag;
2267: abt->alg = alg;
2268: switch (alg) {
2269: case 1: /* alg: "cyclic" */
2270: for (maxRows = 0, i = 0; i < size; i++) maxRows = PetscMax(maxRows, B->rmap->range[i + 1] - B->rmap->range[i]);
2271: bufsiz = A->cmap->N * maxRows;
2272: PetscCall(PetscMalloc2(bufsiz, &abt->buf[0], bufsiz, &abt->buf[1]));
2273: break;
2274: default: /* alg: "allgatherv" */
2275: PetscCall(PetscMalloc2(B->rmap->n * B->cmap->N, &abt->buf[0], B->rmap->N * B->cmap->N, &abt->buf[1]));
2276: PetscCall(PetscMalloc2(size, &abt->recvcounts, size + 1, &abt->recvdispls));
2277: for (i = 0; i <= size; i++) PetscCall(PetscMPIIntCast(B->rmap->range[i] * A->cmap->N, &abt->recvdispls[i]));
2278: for (i = 0; i < size; i++) PetscCall(PetscMPIIntCast(abt->recvdispls[i + 1] - abt->recvdispls[i], &abt->recvcounts[i]));
2279: break;
2280: }
2282: C->product->data = abt;
2283: C->product->destroy = MatProductCtxDestroy_MatMatTransMult_MPIDense_MPIDense;
2284: C->ops->mattransposemultnumeric = MatMatTransposeMultNumeric_MPIDense_MPIDense;
2285: PetscFunctionReturn(PETSC_SUCCESS);
2286: }
2288: static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(Mat A, Mat B, Mat C)
2289: {
2290: Mat_MPIDense *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
2291: MatProductCtx_MatTransMultDense *abt;
2292: MPI_Comm comm;
2293: PetscMPIInt rank, size, sendto, recvfrom, recvisfrom;
2294: PetscScalar *sendbuf, *recvbuf = NULL, *cv;
2295: PetscInt i, cK = A->cmap->N, sendsiz, recvsiz, k, j, bn;
2296: PetscScalar _DOne = 1.0, _DZero = 0.0;
2297: const PetscScalar *av, *bv;
2298: PetscBLASInt cm, cn, ck, alda, blda = 0, clda;
2299: MPI_Request reqs[2];
2300: const PetscInt *ranges;
2302: PetscFunctionBegin;
2303: MatCheckProduct(C, 3);
2304: PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2305: abt = (MatProductCtx_MatTransMultDense *)C->product->data;
2306: PetscCall(PetscObjectGetComm((PetscObject)C, &comm));
2307: PetscCallMPI(MPI_Comm_rank(comm, &rank));
2308: PetscCallMPI(MPI_Comm_size(comm, &size));
2309: PetscCall(MatDenseGetArrayRead(a->A, &av));
2310: PetscCall(MatDenseGetArrayRead(b->A, &bv));
2311: PetscCall(MatDenseGetArrayWrite(c->A, &cv));
2312: PetscCall(MatDenseGetLDA(a->A, &i));
2313: PetscCall(PetscBLASIntCast(i, &alda));
2314: PetscCall(MatDenseGetLDA(b->A, &i));
2315: PetscCall(PetscBLASIntCast(i, &blda));
2316: PetscCall(MatDenseGetLDA(c->A, &i));
2317: PetscCall(PetscBLASIntCast(i, &clda));
2318: PetscCall(MatGetOwnershipRanges(B, &ranges));
2319: bn = B->rmap->n;
2320: if (blda == bn) {
2321: sendbuf = (PetscScalar *)bv;
2322: } else {
2323: sendbuf = abt->buf[0];
2324: for (k = 0, i = 0; i < cK; i++) {
2325: for (j = 0; j < bn; j++, k++) sendbuf[k] = bv[i * blda + j];
2326: }
2327: }
2328: if (size > 1) {
2329: sendto = (rank + size - 1) % size;
2330: recvfrom = (rank + size + 1) % size;
2331: } else {
2332: sendto = recvfrom = 0;
2333: }
2334: PetscCall(PetscBLASIntCast(cK, &ck));
2335: PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm));
2336: recvisfrom = rank;
2337: for (i = 0; i < size; i++) {
2338: /* we have finished receiving in sending, bufs can be read/modified */
2339: PetscMPIInt nextrecvisfrom = (recvisfrom + 1) % size; /* which process the next recvbuf will originate on */
2340: PetscInt nextbn = ranges[nextrecvisfrom + 1] - ranges[nextrecvisfrom];
2342: if (nextrecvisfrom != rank) {
2343: /* start the cyclic sends from sendbuf, to recvbuf (which will switch to sendbuf) */
2344: sendsiz = cK * bn;
2345: recvsiz = cK * nextbn;
2346: recvbuf = (i & 1) ? abt->buf[0] : abt->buf[1];
2347: PetscCallMPI(MPIU_Isend(sendbuf, sendsiz, MPIU_SCALAR, sendto, abt->tag, comm, &reqs[0]));
2348: PetscCallMPI(MPIU_Irecv(recvbuf, recvsiz, MPIU_SCALAR, recvfrom, abt->tag, comm, &reqs[1]));
2349: }
2351: /* local aseq * sendbuf^T */
2352: PetscCall(PetscBLASIntCast(ranges[recvisfrom + 1] - ranges[recvisfrom], &cn));
2353: if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "T", &cm, &cn, &ck, &_DOne, av, &alda, sendbuf, &cn, &_DZero, cv + clda * ranges[recvisfrom], &clda));
2355: if (nextrecvisfrom != rank) {
2356: /* wait for the sends and receives to complete, swap sendbuf and recvbuf */
2357: PetscCallMPI(MPI_Waitall(2, reqs, MPI_STATUSES_IGNORE));
2358: }
2359: bn = nextbn;
2360: recvisfrom = nextrecvisfrom;
2361: sendbuf = recvbuf;
2362: }
2363: PetscCall(MatDenseRestoreArrayRead(a->A, &av));
2364: PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
2365: PetscCall(MatDenseRestoreArrayWrite(c->A, &cv));
2366: PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2367: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2368: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2369: PetscFunctionReturn(PETSC_SUCCESS);
2370: }
2372: static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(Mat A, Mat B, Mat C)
2373: {
2374: Mat_MPIDense *a = (Mat_MPIDense *)A->data, *b = (Mat_MPIDense *)B->data, *c = (Mat_MPIDense *)C->data;
2375: MatProductCtx_MatTransMultDense *abt;
2376: MPI_Comm comm;
2377: PetscMPIInt size, ibn;
2378: PetscScalar *cv, *sendbuf, *recvbuf;
2379: const PetscScalar *av, *bv;
2380: PetscInt blda, i, cK = A->cmap->N, k, j, bn;
2381: PetscScalar _DOne = 1.0, _DZero = 0.0;
2382: PetscBLASInt cm, cn, ck, alda, clda;
2384: PetscFunctionBegin;
2385: MatCheckProduct(C, 3);
2386: PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2387: abt = (MatProductCtx_MatTransMultDense *)C->product->data;
2388: PetscCall(PetscObjectGetComm((PetscObject)A, &comm));
2389: PetscCallMPI(MPI_Comm_size(comm, &size));
2390: PetscCall(MatDenseGetArrayRead(a->A, &av));
2391: PetscCall(MatDenseGetArrayRead(b->A, &bv));
2392: PetscCall(MatDenseGetArrayWrite(c->A, &cv));
2393: PetscCall(MatDenseGetLDA(a->A, &i));
2394: PetscCall(PetscBLASIntCast(i, &alda));
2395: PetscCall(MatDenseGetLDA(b->A, &blda));
2396: PetscCall(MatDenseGetLDA(c->A, &i));
2397: PetscCall(PetscBLASIntCast(i, &clda));
2398: /* copy transpose of B into buf[0] */
2399: bn = B->rmap->n;
2400: sendbuf = abt->buf[0];
2401: recvbuf = abt->buf[1];
2402: for (k = 0, j = 0; j < bn; j++) {
2403: for (i = 0; i < cK; i++, k++) sendbuf[k] = bv[i * blda + j];
2404: }
2405: PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
2406: PetscCall(PetscMPIIntCast(bn * cK, &ibn));
2407: PetscCallMPI(MPI_Allgatherv(sendbuf, ibn, MPIU_SCALAR, recvbuf, abt->recvcounts, abt->recvdispls, MPIU_SCALAR, comm));
2408: PetscCall(PetscBLASIntCast(cK, &ck));
2409: PetscCall(PetscBLASIntCast(c->A->rmap->n, &cm));
2410: PetscCall(PetscBLASIntCast(c->A->cmap->n, &cn));
2411: if (cm && cn && ck) PetscCallBLAS("BLASgemm", BLASgemm_("N", "N", &cm, &cn, &ck, &_DOne, av, &alda, recvbuf, &ck, &_DZero, cv, &clda));
2412: PetscCall(MatDenseRestoreArrayRead(a->A, &av));
2413: PetscCall(MatDenseRestoreArrayRead(b->A, &bv));
2414: PetscCall(MatDenseRestoreArrayWrite(c->A, &cv));
2415: PetscCall(MatSetOption(C, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
2416: PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
2417: PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
2418: PetscFunctionReturn(PETSC_SUCCESS);
2419: }
2421: static PetscErrorCode MatMatTransposeMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2422: {
2423: MatProductCtx_MatTransMultDense *abt;
2425: PetscFunctionBegin;
2426: MatCheckProduct(C, 3);
2427: PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data empty");
2428: abt = (MatProductCtx_MatTransMultDense *)C->product->data;
2429: switch (abt->alg) {
2430: case 1:
2431: PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Cyclic(A, B, C));
2432: break;
2433: default:
2434: PetscCall(MatMatTransposeMultNumeric_MPIDense_MPIDense_Allgatherv(A, B, C));
2435: break;
2436: }
2437: PetscFunctionReturn(PETSC_SUCCESS);
2438: }
2440: static PetscErrorCode MatProductCtxDestroy_MatMatMult_MPIDense_MPIDense(PetscCtxRt data)
2441: {
2442: MatProductCtx_MatMultDense *ab = *(MatProductCtx_MatMultDense **)data;
2444: PetscFunctionBegin;
2445: PetscCall(MatDestroy(&ab->Ce));
2446: PetscCall(MatDestroy(&ab->Ae));
2447: PetscCall(MatDestroy(&ab->Be));
2448: PetscCall(PetscFree(ab));
2449: PetscFunctionReturn(PETSC_SUCCESS);
2450: }
2452: static PetscErrorCode MatMatMultNumeric_MPIDense_MPIDense(Mat A, Mat B, Mat C)
2453: {
2454: MatProductCtx_MatMultDense *ab;
2455: Mat_MPIDense *mdn = (Mat_MPIDense *)A->data;
2456: Mat_MPIDense *b = (Mat_MPIDense *)B->data;
2458: PetscFunctionBegin;
2459: MatCheckProduct(C, 3);
2460: PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Missing product data");
2461: ab = (MatProductCtx_MatMultDense *)C->product->data;
2462: if (ab->Ae && ab->Ce) {
2463: #if PetscDefined(HAVE_ELEMENTAL)
2464: PetscCall(MatConvert_MPIDense_Elemental(A, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Ae));
2465: PetscCall(MatConvert_MPIDense_Elemental(B, MATELEMENTAL, MAT_REUSE_MATRIX, &ab->Be));
2466: PetscCall(MatMatMultNumeric_Elemental(ab->Ae, ab->Be, ab->Ce));
2467: PetscCall(MatConvert(ab->Ce, MATMPIDENSE, MAT_REUSE_MATRIX, &C));
2468: #else
2469: SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined");
2470: #endif
2471: } else {
2472: MPI_Comm comm;
2473: const PetscScalar *read;
2474: PetscScalar *write;
2475: PetscInt lda;
2476: const PetscInt *ranges;
2477: PetscMPIInt size;
2479: if (!mdn->Mvctx) PetscCall(MatSetUpMultiply_MPIDense(A)); /* cannot be done during the symbolic phase because of possible calls to MatProductReplaceMats() */
2480: comm = PetscObjectComm((PetscObject)B);
2481: PetscCallMPI(MPI_Comm_size(comm, &size));
2482: PetscCall(PetscLayoutGetRanges(B->rmap, &ranges));
2483: if (ranges[1] == ranges[size]) {
2484: // optimize for the case where the B matrix is broadcast from rank 0
2485: PetscInt b_lda, be_lda;
2486: Mat b_local = b->A;
2487: Mat b_alloc = NULL;
2488: Mat be_local = ab->Be;
2489: Mat be_alloc = NULL;
2490: PetscMemType b_memtype, be_memtype;
2491: const PetscScalar *b_array = NULL;
2492: MPI_Datatype vector_type;
2493: PetscScalar *be_array = NULL;
2494: PetscMPIInt rank;
2496: PetscCallMPI(MPI_Comm_rank(comm, &rank));
2497: PetscCall(MatDenseGetLDA(be_local, &be_lda));
2498: if (be_lda != B->rmap->N) {
2499: PetscCall(MatDuplicate(be_local, MAT_DO_NOT_COPY_VALUES, &be_alloc));
2500: be_local = be_alloc;
2501: }
2503: if (rank == 0) {
2504: PetscCall(MatDenseGetLDA(b_local, &b_lda));
2505: if (b_lda != B->rmap->N) {
2506: PetscCall(MatDuplicate(b_local, MAT_DO_NOT_COPY_VALUES, &b_alloc));
2507: PetscCall(MatCopy(b_local, b_alloc, DIFFERENT_NONZERO_PATTERN));
2508: b_local = b_alloc;
2509: }
2510: }
2511: vector_type = MPIU_SCALAR;
2512: if (B->cmap->N > 1) {
2513: PetscMPIInt mpi_N;
2515: PetscCall(PetscMPIIntCast(B->cmap->N, &mpi_N));
2516: PetscCallMPI(MPI_Type_contiguous(mpi_N, MPIU_SCALAR, &vector_type));
2517: PetscCallMPI(MPI_Type_commit(&vector_type));
2518: }
2519: PetscCall(MatDenseGetArrayReadAndMemType(b_local, &b_array, &b_memtype));
2520: PetscCall(MatDenseGetArrayWriteAndMemType(be_local, &be_array, &be_memtype));
2521: PetscCall(PetscSFBcastWithMemTypeBegin(mdn->Mvctx, vector_type, b_memtype, b_array, be_memtype, be_array, MPI_REPLACE));
2522: PetscCall(PetscSFBcastEnd(mdn->Mvctx, vector_type, b_array, be_array, MPI_REPLACE));
2523: PetscCall(MatDenseRestoreArrayWriteAndMemType(be_local, &be_array));
2524: PetscCall(MatDenseRestoreArrayReadAndMemType(b_local, &b_array));
2525: if (be_local != ab->Be) PetscCall(MatCopy(be_local, ab->Be, DIFFERENT_NONZERO_PATTERN));
2526: if (B->cmap->N > 1) PetscCallMPI(MPI_Type_free(&vector_type));
2527: PetscCall(MatDestroy(&be_alloc));
2528: PetscCall(MatDestroy(&b_alloc));
2529: } else {
2530: PetscCall(MatDenseGetLDA(B, &lda));
2531: PetscCall(MatDenseGetArrayRead(B, &read));
2532: PetscCall(MatDenseGetArrayWrite(ab->Be, &write));
2533: for (PetscInt i = 0; i < C->cmap->N; ++i) {
2534: PetscCall(PetscSFBcastBegin(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE));
2535: PetscCall(PetscSFBcastEnd(mdn->Mvctx, MPIU_SCALAR, read + i * lda, write + i * ab->Be->rmap->n, MPI_REPLACE));
2536: }
2537: PetscCall(MatDenseRestoreArrayWrite(ab->Be, &write));
2538: PetscCall(MatDenseRestoreArrayRead(B, &read));
2539: }
2540: PetscCall(MatMatMultNumeric_SeqDense_SeqDense(((Mat_MPIDense *)A->data)->A, ab->Be, ((Mat_MPIDense *)C->data)->A));
2541: }
2542: PetscFunctionReturn(PETSC_SUCCESS);
2543: }
2545: static PetscErrorCode MatMatMultSymbolic_MPIDense_MPIDense(Mat A, Mat B, PetscReal fill, Mat C)
2546: {
2547: Mat_Product *product = C->product;
2548: PetscInt alg;
2549: MatProductCtx_MatMultDense *ab;
2550: PetscBool flg;
2552: PetscFunctionBegin;
2553: MatCheckProduct(C, 4);
2554: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Product data not empty");
2555: /* check local size of A and B */
2556: 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 ")",
2557: A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
2559: PetscCall(PetscStrcmp(product->alg, "petsc", &flg));
2560: alg = flg ? 0 : 1;
2562: /* setup C */
2563: PetscCall(MatSetSizes(C, A->rmap->n, B->cmap->n, A->rmap->N, B->cmap->N));
2564: PetscCall(MatSetType(C, MATMPIDENSE));
2565: /* keep the VecType of A, e.g. VECKOKKOS from MatCreateDenseFromVecType(), but only if A is a MATMPIDENSE like C, since the VecType of a MATMPIDENSECUDA or MATMPIDENSEHIP A, e.g. VECCUDA, is not valid for C */
2566: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATMPIDENSE, &flg));
2567: if (flg) PetscCall(MatSetVecType(C, A->defaultvectype));
2568: PetscCall(MatSetUp(C));
2570: /* create data structure for reuse Cdense */
2571: PetscCall(PetscNew(&ab));
2573: switch (alg) {
2574: case 1: /* alg: "elemental" */
2575: #if PetscDefined(HAVE_ELEMENTAL)
2576: /* create elemental matrices Ae and Be */
2577: PetscCall(MatCreate(PetscObjectComm((PetscObject)A), &ab->Ae));
2578: PetscCall(MatSetSizes(ab->Ae, PETSC_DECIDE, PETSC_DECIDE, A->rmap->N, A->cmap->N));
2579: PetscCall(MatSetType(ab->Ae, MATELEMENTAL));
2580: PetscCall(MatSetUp(ab->Ae));
2581: PetscCall(MatSetOption(ab->Ae, MAT_ROW_ORIENTED, PETSC_FALSE));
2583: PetscCall(MatCreate(PetscObjectComm((PetscObject)B), &ab->Be));
2584: PetscCall(MatSetSizes(ab->Be, PETSC_DECIDE, PETSC_DECIDE, B->rmap->N, B->cmap->N));
2585: PetscCall(MatSetType(ab->Be, MATELEMENTAL));
2586: PetscCall(MatSetUp(ab->Be));
2587: PetscCall(MatSetOption(ab->Be, MAT_ROW_ORIENTED, PETSC_FALSE));
2589: /* compute symbolic Ce = Ae*Be */
2590: PetscCall(MatCreate(PetscObjectComm((PetscObject)C), &ab->Ce));
2591: PetscCall(MatMatMultSymbolic_Elemental(ab->Ae, ab->Be, fill, ab->Ce));
2592: #else
2593: SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "PETSC_HAVE_ELEMENTAL not defined");
2594: #endif
2595: break;
2596: default: /* alg: "petsc" */
2597: ab->Ae = NULL;
2598: PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, A->cmap->N, B->cmap->N, NULL, &ab->Be));
2599: ab->Ce = NULL;
2600: break;
2601: }
2603: C->product->data = ab;
2604: C->product->destroy = MatProductCtxDestroy_MatMatMult_MPIDense_MPIDense;
2605: C->ops->matmultnumeric = MatMatMultNumeric_MPIDense_MPIDense;
2606: PetscFunctionReturn(PETSC_SUCCESS);
2607: }
2609: static PetscErrorCode MatProductSetFromOptions_MPIDense_AB(Mat C)
2610: {
2611: Mat_Product *product = C->product;
2612: const char *algTypes[2] = {"petsc", "elemental"};
2613: PetscInt alg, nalg = PetscDefined(HAVE_ELEMENTAL) ? 2 : 1;
2614: PetscBool flg = PETSC_FALSE;
2616: PetscFunctionBegin;
2617: /* Set default algorithm */
2618: alg = 0; /* default is PETSc */
2619: PetscCall(PetscStrcmp(product->alg, "default", &flg));
2620: if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2622: /* Get runtime option */
2623: PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_AB", "Mat");
2624: PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_AB", algTypes, nalg, algTypes[alg], &alg, &flg));
2625: PetscOptionsEnd();
2626: if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2628: C->ops->matmultsymbolic = MatMatMultSymbolic_MPIDense_MPIDense;
2629: C->ops->productsymbolic = MatProductSymbolic_AB;
2630: PetscFunctionReturn(PETSC_SUCCESS);
2631: }
2633: static PetscErrorCode MatProductSetFromOptions_MPIDense_AtB(Mat C)
2634: {
2635: Mat_Product *product = C->product;
2636: Mat A = product->A, B = product->B;
2638: PetscFunctionBegin;
2639: 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 ")",
2640: A->rmap->rstart, A->rmap->rend, B->rmap->rstart, B->rmap->rend);
2641: C->ops->transposematmultsymbolic = MatTransposeMatMultSymbolic_MPIDense_MPIDense;
2642: C->ops->productsymbolic = MatProductSymbolic_AtB;
2643: PetscFunctionReturn(PETSC_SUCCESS);
2644: }
2646: static PetscErrorCode MatProductSetFromOptions_MPIDense_ABt(Mat C)
2647: {
2648: Mat_Product *product = C->product;
2649: const char *algTypes[2] = {"allgatherv", "cyclic"};
2650: PetscInt alg, nalg = 2;
2651: PetscBool flg = PETSC_FALSE;
2653: PetscFunctionBegin;
2654: /* Set default algorithm */
2655: alg = 0; /* default is allgatherv */
2656: PetscCall(PetscStrcmp(product->alg, "default", &flg));
2657: if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2659: /* Get runtime option */
2660: if (product->api_user) {
2661: PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatMatTransposeMult", "Mat");
2662: PetscCall(PetscOptionsEList("-matmattransmult_mpidense_mpidense_via", "Algorithmic approach", "MatMatTransposeMult", algTypes, nalg, algTypes[alg], &alg, &flg));
2663: PetscOptionsEnd();
2664: } else {
2665: PetscOptionsBegin(PetscObjectComm((PetscObject)C), ((PetscObject)C)->prefix, "MatProduct_ABt", "Mat");
2666: PetscCall(PetscOptionsEList("-mat_product_algorithm", "Algorithmic approach", "MatProduct_ABt", algTypes, nalg, algTypes[alg], &alg, &flg));
2667: PetscOptionsEnd();
2668: }
2669: if (flg) PetscCall(MatProductSetAlgorithm(C, algTypes[alg]));
2671: C->ops->mattransposemultsymbolic = MatMatTransposeMultSymbolic_MPIDense_MPIDense;
2672: C->ops->productsymbolic = MatProductSymbolic_ABt;
2673: PetscFunctionReturn(PETSC_SUCCESS);
2674: }
2676: static PetscErrorCode MatProductSetFromOptions_MPIDense(Mat C)
2677: {
2678: Mat_Product *product = C->product;
2680: PetscFunctionBegin;
2681: switch (product->type) {
2682: case MATPRODUCT_AB:
2683: PetscCall(MatProductSetFromOptions_MPIDense_AB(C));
2684: break;
2685: case MATPRODUCT_AtB:
2686: PetscCall(MatProductSetFromOptions_MPIDense_AtB(C));
2687: break;
2688: case MATPRODUCT_ABt:
2689: PetscCall(MatProductSetFromOptions_MPIDense_ABt(C));
2690: break;
2691: default:
2692: break;
2693: }
2694: PetscFunctionReturn(PETSC_SUCCESS);
2695: }
2697: PetscErrorCode MatDenseScatter_Private(PetscSF sf, Mat X, Mat Y, InsertMode mode, ScatterMode smode)
2698: {
2699: const PetscScalar *in;
2700: PetscScalar *out;
2701: PetscSF vsf;
2702: PetscInt N, ny, rld, lld;
2703: PetscMemType mtype[2];
2704: MPI_Op op = MPI_OP_NULL;
2706: PetscFunctionBegin;
2710: if (mode == INSERT_VALUES) op = MPI_REPLACE;
2711: else if (mode == ADD_VALUES) op = MPIU_SUM;
2712: else if (mode == MAX_VALUES) op = MPIU_MAX;
2713: else if (mode == MIN_VALUES) op = MPIU_MIN;
2714: PetscCheck(op != MPI_OP_NULL, PetscObjectComm((PetscObject)sf), PETSC_ERR_SUP, "Unsupported InsertMode %d in MatDenseScatter_Private()", mode);
2715: PetscCheck(smode == SCATTER_FORWARD || smode == SCATTER_REVERSE, PetscObjectComm((PetscObject)sf), PETSC_ERR_SUP, "Unsupported ScatterMode %d in MatDenseScatter_Private()", smode);
2716: PetscCall(MatGetSize(X, NULL, &N));
2717: PetscCall(MatGetSize(Y, NULL, &ny));
2718: PetscCheck(N == ny, PetscObjectComm((PetscObject)sf), PETSC_ERR_ARG_SIZ, "Matrix column sizes must match: %" PetscInt_FMT " != %" PetscInt_FMT, N, ny);
2719: PetscCall(MatDenseGetLDA(X, &rld));
2720: PetscCall(MatDenseGetLDA(Y, &lld));
2721: /* get cached or create new strided PetscSF when the number of columns is greater than one */
2722: if (N > 1) {
2723: PetscCall(PetscObjectQuery((PetscObject)sf, "_MatDenseScatter_StridedSF", (PetscObject *)&vsf));
2724: if (vsf) {
2725: PetscInt nr[2], nl[2];
2727: PetscCall(PetscSFGetGraph(sf, nr, nl, NULL, NULL));
2728: PetscCall(PetscSFGetGraph(vsf, nr + 1, nl + 1, NULL, NULL));
2729: if (N * nr[0] != nr[1] || N * nl[0] != nl[1]) vsf = NULL;
2730: }
2731: if (!vsf) {
2732: PetscCall(PetscSFCreateStridedSF(sf, N, rld, lld, &vsf));
2733: PetscCall(PetscObjectCompose((PetscObject)sf, "_MatDenseScatter_StridedSF", (PetscObject)vsf));
2734: PetscCall(PetscObjectDereference((PetscObject)vsf));
2735: }
2736: } else vsf = sf;
2737: /* the output array is accessed in read and write mode,
2738: but write-only in the INSERT_VALUES case could be worth exploring */
2739: PetscCall(MatDenseGetArrayReadAndMemType(X, &in, &mtype[0]));
2740: PetscCall(MatDenseGetArrayAndMemType(Y, &out, &mtype[1]));
2741: if (smode == SCATTER_FORWARD) {
2742: PetscCall(PetscSFBcastWithMemTypeBegin(vsf, vsf->vscat.unit, mtype[0], in, mtype[1], out, op));
2743: PetscCall(PetscSFBcastEnd(vsf, vsf->vscat.unit, in, out, op));
2744: } else {
2745: PetscCall(PetscSFReduceWithMemTypeBegin(vsf, vsf->vscat.unit, mtype[0], in, mtype[1], out, op));
2746: PetscCall(PetscSFReduceEnd(vsf, vsf->vscat.unit, in, out, op));
2747: }
2748: PetscCall(MatDenseRestoreArrayAndMemType(Y, &out));
2749: PetscCall(MatDenseRestoreArrayReadAndMemType(X, &in));
2750: PetscFunctionReturn(PETSC_SUCCESS);
2751: }