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